[pypy-commit] pypy default: Use OneDimIterator if we have one already. Fix it since tests started failing,

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r49941:8e56ec1b31a3
Date: 2011-11-29 10:10 +0200
http://bitbucket.org/pypy/pypy/changeset/8e56ec1b31a3/

Log:	Use OneDimIterator if we have one already. Fix it since tests
	started failing, make test_zjit look better again.

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
@@ -180,10 +180,10 @@
         return self.offset
 
 class OneDimIterator(BaseIterator):
-    def __init__(self, start, step, size):
+    def __init__(self, start, step, stop):
         self.offset = start
         self.step = step
-        self.size = size
+        self.size = stop * step + start
 
     def next(self, shapelen):
         arr = instantiate(OneDimIterator)
@@ -193,7 +193,7 @@
         return arr
 
     def done(self):
-        return self.offset >= self.size
+        return self.offset == self.size
 
     def get_offset(self):
         return self.offset
@@ -1109,8 +1109,8 @@
     def start_iter(self, res_shape=None):
         if res_shape is not None and res_shape != self.shape:
             return BroadcastIterator(self, res_shape)
-        # XXX there is a possible optimization here with SingleDimViewIterator
-        #     ignore for now
+        if len(self.shape) == 1:
+            return OneDimIterator(self.start, self.strides[0], self.shape[0])
         return ViewIterator(self)
 
     def setitem(self, item, value):
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
@@ -246,12 +246,12 @@
     def test_slice(self):
         result = self.run("slice")
         assert result == 18
-        # arraylen_gc are removed by the backend, would be good if they weren't
-        # here though
-        self.check_simple_loop({'int_mul': 2, 'getarrayitem_raw': 2, 'float_add': 1,
-                          'setarrayitem_raw': 1, 'int_add': 3,
-                          'int_lt': 1, 'guard_true': 1, 'jump': 1,
-                          'arraylen_gc': 4})
+        self.check_simple_loop({'getarrayitem_raw': 2,
+                                'float_add': 1,
+                                'setarrayitem_raw': 1,
+                                'int_add': 3,
+                                'int_ge': 1, 'guard_false': 1,
+                                'jump': 1})
 
     def define_slice2():
         return """
@@ -265,12 +265,9 @@
     def test_slice2(self):
         result = self.run("slice2")
         assert result == 15
-        # arraylen_gc are removed by the backend, would be good if they weren't
-        # here though
-        self.check_simple_loop({'int_mul': 2, 'getarrayitem_raw': 2, 'float_add': 1,
-                                'setarrayitem_raw': 1, 'int_add': 1,
-                                'int_lt': 1, 'guard_true': 1, 'jump': 1,
-                                'arraylen_gc': 4})
+        self.check_simple_loop({'getarrayitem_raw': 2, 'float_add': 1,
+                                'setarrayitem_raw': 1, 'int_add': 3,
+                                'int_ge': 1, 'guard_false': 1, 'jump': 1})
 
     def define_multidim():
         return """