[pypy-commit] pypy numppy-flatitter: fix for z_jit test, now passes

[mattip]

Author: mattip
Branch: numppy-flatitter
Changeset: r51839:5a2c53898154
Date: 2012-01-19 11:30 +0200
http://bitbucket.org/pypy/pypy/changeset/5a2c53898154/

Log:	fix for z_jit test, now passes

diff --git a/pypy/module/micronumpy/compile.py b/pypy/module/micronumpy/compile.py
--- a/pypy/module/micronumpy/compile.py
+++ b/pypy/module/micronumpy/compile.py
@@ -9,7 +9,7 @@
 from pypy.module.micronumpy import interp_boxes
 from pypy.module.micronumpy.interp_dtype import get_dtype_cache
 from pypy.module.micronumpy.interp_numarray import (Scalar, BaseArray,
-     scalar_w, W_NDimArray, array)
+     scalar_w, W_NDimArray, array, W_FlatIterator)
 from pypy.module.micronumpy import interp_ufuncs
 from pypy.rlib.objectmodel import specialize, instantiate
 
@@ -313,6 +313,21 @@
     def __repr__(self):
         return 'Range(%s)' % self.v
 
+class FlatIterFromRange(Node):
+    def __init__(self, v):
+        self.v = int(v)
+
+    def execute(self, interp):
+        w_list = interp.space.newlist(
+            [interp.space.wrap(float(i)) for i in range(self.v)]
+        )
+        dtype = get_dtype_cache(interp.space).w_float64dtype
+        return W_FlatIterator(array(interp.space, w_list, w_dtype=dtype, 
+                                w_order=None))
+
+    def __repr__(self):
+        return 'FlatFromRange(%s)' % self.v
+
 class Code(Node):
     def __init__(self, statements):
         self.statements = statements
@@ -422,6 +437,7 @@
     ('=', 'assign'),
     (',', 'comma'),
     ('\|', 'pipe'),
+    ('\!', 'exclaim'),
     ('\(', 'paren_left'),
     ('\)', 'paren_right'),
 ]
@@ -528,6 +544,10 @@
                 stack.append(RangeConstant(tokens.pop().v))
                 end = tokens.pop()
                 assert end.name == 'pipe'
+            elif token.name == 'exclaim':
+                stack.append(FlatIterFromRange(tokens.pop().v))
+                end = tokens.pop()
+                assert end.name == 'exclaim'
             elif accept_comma and token.name == 'comma':
                 continue
             else:
diff --git a/pypy/module/micronumpy/interp_iter.py b/pypy/module/micronumpy/interp_iter.py
--- a/pypy/module/micronumpy/interp_iter.py
+++ b/pypy/module/micronumpy/interp_iter.py
@@ -137,7 +137,7 @@
         res._done = done
         return res
 
-    #@jit.unroll_safe
+    @jit.unroll_safe
     def next_skip_x(self, shapelen, step):
         shapelen = jit.promote(len(self.res_shape))
         offset = self.offset
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
@@ -50,8 +50,6 @@
             if not len(interp.results):
                 raise Exception("need results")
             w_res = interp.results[-1]
-            if isinstance(w_res, W_FlatIterator):
-                w_res = w_res.next()
             if isinstance(w_res, BaseArray):
                 concr = w_res.get_concrete_or_scalar()
                 sig = concr.find_sig()
@@ -373,8 +371,7 @@
                                 'arraylen_gc': 1})
     def define_flat_iter():
         return '''
-        a = |30|
-        a -> flat
+        !30!
         '''
 
     def test_flat_iter(self):
@@ -383,16 +380,14 @@
 
     def define_flat_getitem():
         return '''
-        a = |30|
-        b = a -> flat
-        b -> 6
+        a = !30!
+        a -> 6
         '''
 
     def test_flat_getitem(self):
         result = self.run("flat_getitem")
-        assert result == 4
-        self.check_trace_count(1)
-        self.check_simple_loop({})
+        assert result == 6.0
+        self.check_trace_count(0)
 
 class TestNumpyOld(LLJitMixin):
     def setup_class(cls):