[pypy-commit] pypy numpy-singledim: micronumpy: second version of setslice including jit driver and test

[justinpeel]

Author: Justin Peel <notmuchtotell at gmail.com>
Branch: numpy-singledim
Changeset: r45644:182a6ac8d391
Date: 2011-07-14 18:53 -0600
http://bitbucket.org/pypy/pypy/changeset/182a6ac8d391/

Log:	micronumpy: second version of setslice including jit driver and test

diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -21,6 +21,8 @@
                              reds = ['result_size', 'i', 'self', 'result'])
 all_driver = jit.JitDriver(greens=['signature'], reds=['i', 'size', 'self'])
 any_driver = jit.JitDriver(greens=['signature'], reds=['i', 'size', 'self'])
+slice_driver1 = jit.JitDriver(greens=['signature'], reds=['i', 'j', 'step', 'stop', 'self', 'arr'])
+slice_driver2 = jit.JitDriver(greens=['signature'], reds=['i', 'j', 'step', 'stop', 'self', 'arr'])
 
 class Signature(object):
     def __init__(self):
@@ -257,7 +259,7 @@
         return self.get_concrete().descr_str(space)
 
     def descr_getitem(self, space, w_idx):
-        # TODO: indexing by tuples
+        # TODO: indexing by tuples and lists
         start, stop, step, slice_length = space.decode_index4(w_idx, self.find_size())
         if step == 0:
             # Single index
@@ -268,6 +270,7 @@
             return space.wrap(res)
 
     def descr_setitem(self, space, w_idx, w_value):
+        # TODO: indexing by tuples and lists
         self.invalidated()
         start, stop, step, slice_length = space.decode_index4(w_idx,
                                                               self.find_size())
@@ -513,53 +516,50 @@
     def descr_str(self,space):
         return space.wrap("[" + " ".join(self._getnums(True)) + "]")
 
+    @unwrap_spec(item=int, value=float)
     def descr_setitem(self, space, item, value):
         item = self.getindex(space, item)
         self.invalidated()
         self.storage[item] = value
 
-    def descr_setslice(self, space, start, stop, step, slice_length, w_value):
+    def _setslice1(self, start, stop, step, arr):
+        signature = Signature()
+        new_sig = self.signature.transition(signature)
         i = start
+        j = 0
+        while i < stop:
+            slice_driver1.jit_merge_point(signature=signature, self=self,
+                    step=step, stop=stop, i=i, j=j, arr=arr)
+            self.storage[i] = arr.eval(j)
+            j += 1
+            i += step
+
+    def _setslice2(self, start, stop, step, arr):
+        signature = Signature()
+        new_sig = self.signature.transition(signature)
+        i = start
+        j = 0
+        while i > stop:
+            slice_driver2.jit_merge_point(signature=signature, self=self,
+                    step=step, stop=stop, i=i, j=j, arr=arr)
+            self.storage[i] = arr.eval(j)
+            j += 1
+            i += step
+
+    def descr_setslice(self, space, start, stop, step, slice_length, arr):
+        i = start
+        if stop < 0:
+            stop += self.find_size()
         if step > 0:
             stop = min(stop, self.find_size())
         else:
             stop = max(stop, 0)
-        if isinstance(w_value, BaseArray):
-            j = 0
-            if step > 0:
-                while i < stop:
-                    self.storage[i] = w_value.get_concrete().getitem(j)
-                    i += step
-                    j += 1
-            else:
-                while i > stop:
-                    self.storage[i] = w_value.get_concrete().getitem(j)
-                    i += step
-                    j += 1
-        elif space.issequence_w(w_value):
-            l = space.listview(w_value)
-            if step > 0:
-                for w_elem in l:
-                    self.storage[i] = space.float_w(space.float(w_elem))
-                    i += step
-                    if i >= stop:
-                        break
-            else:
-                for w_elem in l:
-                    self.storage[i] = space.float_w(space.float(w_elem))
-                    i += step
-                    if i <= stop:
-                        break
+        if not isinstance(arr, BaseArray):
+            arr = convert_to_array(space, arr)
+        if step > 0:
+            self._setslice1(start, stop, step, arr)
         else:
-            value = space.float_w(space.float(w_value))
-            if step > 0:
-                while i < stop:
-                    self.storage[i] = value
-                    i += step
-            else:
-                while i > stop:
-                    self.storage[i] = value
-                    i += step
+            self._setslice2(start, stop, step, arr)
 
     def __del__(self):
         lltype.free(self.storage, flavor='raw')
diff --git a/pypy/module/micronumpy/test/test_zjit.py b/pypy/module/micronumpy/test/test_zjit.py
--- a/pypy/module/micronumpy/test/test_zjit.py
+++ b/pypy/module/micronumpy/test/test_zjit.py
@@ -5,6 +5,7 @@
 from pypy.module.micronumpy.interp_ufuncs import negative
 from pypy.module.micronumpy.compile import numpy_compile
 from pypy.rlib.objectmodel import specialize
+from pypy.rlib.nonconst import NonConstant
 
 class FakeSpace(object):
     w_ValueError = None
@@ -248,6 +249,21 @@
                           'int_lt': 1, 'guard_true': 1, 'jump': 1})
         assert result == f(5)
 
+    def test_setslice(self):
+        space = self.space
+
+        def f(i):
+            step = NonConstant(3)
+            ar = SingleDimArray(step*i)
+            ar2 = SingleDimArray(i)
+            ar.descr_setslice(space, 0, step*i, step, i, ar2)
+            return ar.get_concrete().storage[3]
+
+        result = self.meta_interp(f, [5], listops=True, backendopt=True)
+        self.check_loops({'getarrayitem_raw': 1,
+                          'setarrayitem_raw': 1, 'int_add': 2,
+                          'int_lt': 1, 'guard_true': 1, 'jump': 1})
+
 class TestTranslation(object):
     def test_compile(self):
         x = numpy_compile('aa+f*f/a-', 10)