[pypy-svn] pypy numpy-exp: Make auto-lazy-forcing not break semantics.

[alex_gaynor]

Author: Alex Gaynor <alex.gaynor at gmail.com>
Branch: numpy-exp
Changeset: r43899:1486788fda44
Date: 2011-05-04 19:03 -0400
http://bitbucket.org/pypy/pypy/changeset/1486788fda44/

Log:	Make auto-lazy-forcing not break semantics.

diff --git a/pypy/module/micronumpy/numarray.py b/pypy/module/micronumpy/numarray.py
--- a/pypy/module/micronumpy/numarray.py
+++ b/pypy/module/micronumpy/numarray.py
@@ -145,6 +145,9 @@
 JITCODES = {}
 
 class BaseArray(Wrappable):
+    def __init__(self):
+        self.invalidates = []
+
     def force(self):
         code = self.compile()
         try:
@@ -156,16 +159,24 @@
         # (we still have to compile new bytecode, but too bad)
         return compute(code)
 
+    def invalidated(self):
+        for arr in self.invalidates:
+            arr.force_if_needed()
+        self.invalidates = []
+
     def _binop_impl(bytecode):
         def impl(self, space, w_other):
             if isinstance(w_other, BaseArray):
-                return space.wrap(BinOp(bytecode, self, w_other))
+                res = space.wrap(BinOp(bytecode, self, w_other))
+                w_other.invalidates.append(res)
             else:
-                return space.wrap(BinOp(
+                res = space.wrap(BinOp(
                     bytecode,
                     self,
                     FloatWrapper(space.float_w(w_other))
                 ))
+            self.invalidates.append(res)
+            return res
         return func_with_new_name(impl, "binop_%s_impl" % bytecode)
 
     descr_add = _binop_impl("a")
@@ -182,6 +193,7 @@
     """
 
     def __init__(self, float_value):
+        BaseArray.__init__(self)
         self.float_value = float_value
 
     def compile(self):
@@ -193,18 +205,39 @@
     """
 
     def __init__(self, opcode, left, right):
+        BaseArray.__init__(self)
         self.opcode = opcode
         self.left = left
         self.right = right
 
+        self.forced_result = None
+
     def compile(self):
         left_code = self.left.compile()
         right_code = self.right.compile()
         return left_code.merge(self.opcode, right_code)
 
-BaseArray.typedef = TypeDef(
+    def force_if_needed(self):
+        if self.forced_result is None:
+            self.forced_result = self.force()
+
+    @unwrap_spec(item=int)
+    def descr_getitem(self, space, item):
+        self.force_if_needed()
+        return self.forced_result.descr_getitem(space, item)
+
+    @unwrap_spec(item=int, value=float)
+    def descr_setitem(self, space, item, value):
+        self.forced_if_needed()
+        self.invalidated()
+        return self.forced_result.descr_setitem(space, item, value)
+
+
+BinOp.typedef = TypeDef(
     'Operation',
-    force = interp2app(BaseArray.force),
+    __getitem__ = interp2app(BinOp.descr_getitem),
+    __setitem__ = interp2app(BinOp.descr_setitem),
+
     __add__ = interp2app(BaseArray.descr_add),
     __sub__ = interp2app(BaseArray.descr_sub),
     __mul__ = interp2app(BaseArray.descr_mul),
@@ -213,6 +246,7 @@
 
 class SingleDimArray(BaseArray):
     def __init__(self, size):
+        BaseArray.__init__(self)
         self.size = size
         self.storage = lltype.malloc(TP, size, zero=True,
                                      flavor='raw', track_allocation=False)
@@ -239,11 +273,9 @@
         if item > self.size:
             raise operationerrfmt(space.w_TypeError,
               '%d above array size', item)
+        self.invalidated()
         self.storage[item] = value
 
-    def force(self):
-        return self
-
     def __del__(self):
         lltype.free(self.storage, flavor='raw')
 
@@ -270,5 +302,4 @@
     __sub__ = interp2app(BaseArray.descr_sub),
     __mul__ = interp2app(BaseArray.descr_mul),
     __div__ = interp2app(BaseArray.descr_div),
-    force = interp2app(SingleDimArray.force),
 )
\ No newline at end of file
diff --git a/pypy/module/micronumpy/test/test_numpy.py b/pypy/module/micronumpy/test/test_numpy.py
--- a/pypy/module/micronumpy/test/test_numpy.py
+++ b/pypy/module/micronumpy/test/test_numpy.py
@@ -23,7 +23,6 @@
         from numpy import array
         a = array(range(5))
         b = a + a
-        b = b.force()
         for i in range(5):
             assert b[i] == i + i
 
@@ -32,7 +31,6 @@
         a = array(range(5))
         b = array(reversed(range(5)))
         c = a + b
-        c = c.force()
         for i in range(5):
             assert c[i] == 4
 
@@ -40,14 +38,13 @@
         from numpy import array
         a = array(range(5))
         b = a + 5
-        b = b.force()
         for i in range(5):
             assert b[i] == i + 5
 
     def test_subtract(self):
         from numpy import array
         a = array(range(5))
-        b = (a - a).force()
+        b = a - a
         for i in range(5):
             assert b[i] == 0
 
@@ -55,14 +52,14 @@
         from numpy import array
         a = array(range(5))
         b = array([1, 1, 1, 1, 1])
-        c = (a - b).force()
+        c = a - b
         for i in range(5):
             assert c[i] == i - 1
 
     def test_subtract_constant(self):
         from numpy import array
         a = array(range(5))
-        b = (a - 5).force()
+        b = a - 5
         for i in range(5):
             assert b[i] == i - 5
 
@@ -70,7 +67,6 @@
         from numpy import array
         a = array(range(5))
         b = a * a
-        b = b.force()
         for i in range(5):
             assert b[i] == i * i
 
@@ -78,14 +74,13 @@
         from numpy import array
         a = array(range(5))
         b = a * 5
-        b = b.force()
         for i in range(5):
             assert b[i] == i * 5
 
     def test_div(self):
         from numpy import array
         a = array(range(1, 6))
-        b = (a / a).force()
+        b = a / a
         for i in range(5):
             assert b[i] == 1
 
@@ -93,17 +88,25 @@
         from numpy import array
         a = array(range(5))
         b = array([2, 2, 2, 2, 2])
-        c = (a / b).force()
+        c = a / b
         for i in range(5):
             assert c[i] == i / 2.0
 
     def test_div_constant(self):
         from numpy import array
         a = array(range(5))
-        b = (a / 5.0).force()
+        b = a / 5.0
         for i in range(5):
             assert b[i] == i / 5.0
 
+    def test_auto_force(self):
+        from numpy import array
+        a = array(range(5))
+        b = a - 1
+        a[2] = 3
+        for i in range(5):
+            assert b[i] == i - 1
+
 class AppTestNumpy(object):
     def setup_class(cls):
         py.test.skip("unimplemented")