[pypy-commit] pypy translation-cleanup: Simplify HostCode._initialize()

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: translation-cleanup
Changeset: r57958:6f2151c95eb7
Date: 2012-10-05 23:20 +0100
http://bitbucket.org/pypy/pypy/changeset/6f2151c95eb7/

Log:	Simplify HostCode._initialize()

	* Use HostCode.formalargcount instead of recomputing it.
	* Flowspacify FSFrame.save_locals_stack(), .restore_locals_stack()
	and .init_cells()
	* Add test for the corner case of function arguments that are also
	cell variables

diff --git a/pypy/objspace/flow/bytecode.py b/pypy/objspace/flow/bytecode.py
--- a/pypy/objspace/flow/bytecode.py
+++ b/pypy/objspace/flow/bytecode.py
@@ -42,26 +42,18 @@
 
     def _initialize(self):
         # Precompute what arguments need to be copied into cellvars
-        self._args_as_cellvars = []
+        self._args_as_cellvars = {}
 
         if self.co_cellvars:
-            argcount = self.co_argcount
-            assert argcount >= 0     # annotator hint
-            if self.co_flags & CO_VARARGS:
-                argcount += 1
-            if self.co_flags & CO_VARKEYWORDS:
-                argcount += 1
             # Cell vars could shadow already-set arguments.
             # See comment in PyCode._initialize()
             argvars  = self.co_varnames
             cellvars = self.co_cellvars
             for i in range(len(cellvars)):
                 cellname = cellvars[i]
-                for j in range(argcount):
+                for j in range(self.formalargcount):
                     if cellname == argvars[j]:
                         # argument j has the same name as the cell var i
-                        while len(self._args_as_cellvars) <= i:
-                            self._args_as_cellvars.append(-1)   # pad
                         self._args_as_cellvars[i] = j
 
     @classmethod
diff --git a/pypy/objspace/flow/flowcontext.py b/pypy/objspace/flow/flowcontext.py
--- a/pypy/objspace/flow/flowcontext.py
+++ b/pypy/objspace/flow/flowcontext.py
@@ -247,6 +247,21 @@
         self.valuestackdepth = code.co_nlocals
         self.locals_stack_w = [None] * (code.co_stacksize + code.co_nlocals)
 
+    def save_locals_stack(self):
+        return self.locals_stack_w[:self.valuestackdepth]
+
+    def restore_locals_stack(self, items_w):
+        self.locals_stack_w[:len(items_w)] = items_w
+        self.init_cells()
+        self.dropvaluesuntil(len(items_w))
+
+    def init_cells(self):
+        if self.cells is None:
+            return
+        args_to_copy = self.pycode._args_as_cellvars
+        for cellnum, argnum in args_to_copy.iteritems():
+            self.cells[cellnum].set(self.locals_stack_w[argnum])
+
     def getstate(self):
         # getfastscope() can return real None, for undefined locals
         data = self.save_locals_stack()
diff --git a/pypy/objspace/flow/test/test_objspace.py b/pypy/objspace/flow/test/test_objspace.py
--- a/pypy/objspace/flow/test/test_objspace.py
+++ b/pypy/objspace/flow/test/test_objspace.py
@@ -1054,7 +1054,7 @@
         with py.test.raises(FlowingError):
             self.codetest(f)
 
-    def test_cellvar(self):
+    def test_cellvar_store(self):
         def f():
             x = 5
             return x
@@ -1063,6 +1063,18 @@
         assert len(graph.startblock.exits) == 1
         assert graph.startblock.exits[0].target == graph.returnblock
 
+    def test_arg_as_cellvar(self):
+        def f(x, y, z):
+            a, b, c = 1, 2, 3
+            z = b
+            return z
+            lambda: (a, b, x, z) # make cell variables
+        graph = self.codetest(f)
+        assert len(graph.startblock.exits) == 1
+        assert graph.startblock.exits[0].target == graph.returnblock
+        assert not graph.startblock.operations
+        assert graph.startblock.exits[0].args[0].value == 2
+
     def test_lambda(self):
         def f():
             g = lambda m, n: n*m