Speed up with statements by storing the __exit__ method on the stack instead of in a temp variable (bumps the magic number for pyc files)

6 files changed, 77 insertions, 55 deletions

diff --git a/Doc/library/dis.rst b/Doc/library/dis.rst
index c03feb7..d83ac3d 100644
--- a/Doc/library/dis.rst
+++ b/Doc/library/dis.rst
@@ -519,21 +519,24 @@ Miscellaneous opcodes.
 
 .. opcode:: WITH_CLEANUP ()
 
-   Cleans up the stack when a :keyword:`with` statement block exits.  TOS is the
-   context manager's :meth:`__exit__` bound method.  Below that are 1--3 values
-   indicating how/why the finally clause was entered:
+   Cleans up the stack when a :keyword:`with` statement block exits.  On top of
+   the stack are 1--3 values indicating how/why the finally clause was entered:
 
-   * SECOND = ``None``
-   * (SECOND, THIRD) = (``WHY_{RETURN,CONTINUE}``), retval
-   * SECOND = ``WHY_*``; no retval below it
-   * (SECOND, THIRD, FOURTH) = exc_info()
+   * TOP = ``None``
+   * (TOP, SECOND) = (``WHY_{RETURN,CONTINUE}``), retval
+   * TOP = ``WHY_*``; no retval below it
+   * (TOP, SECOND, THIRD) = exc_info()
 
-   In the last case, ``TOS(SECOND, THIRD, FOURTH)`` is called, otherwise
-   ``TOS(None, None, None)``.
+   Under them is EXIT, the context manager's :meth:`__exit__` bound method.
 
-   In addition, if the stack represents an exception, *and* the function call
-   returns a 'true' value, this information is "zapped", to prevent ``END_FINALLY``
-   from re-raising the exception.  (But non-local gotos should still be resumed.)
+   In the last case, ``EXIT(TOP, SECOND, THIRD)`` is called, otherwise
+   ``EXIT(None, None, None)``.
+
+   EXIT is removed from the stack, leaving the values above it in the same
+   order. In addition, if the stack represents an exception, *and* the function
+   call returns a 'true' value, this information is "zapped", to prevent
+   ``END_FINALLY`` from re-raising the exception.  (But non-local gotos should
+   still be resumed.)
 
    .. XXX explain the WHY stuff!
 
diff --git a/Lib/compiler/pycodegen.py b/Lib/compiler/pycodegen.py
index b46b1c1..5d227b8 100644
--- a/Lib/compiler/pycodegen.py
+++ b/Lib/compiler/pycodegen.py
@@ -822,14 +822,13 @@ class CodeGenerator:
     def visitWith(self, node):
         body = self.newBlock()
         final = self.newBlock()
-        exitvar = "$exit%d" % self.__with_count
         valuevar = "$value%d" % self.__with_count
         self.__with_count += 1
         self.set_lineno(node)
         self.visit(node.expr)
         self.emit('DUP_TOP')
         self.emit('LOAD_ATTR', '__exit__')
-        self._implicitNameOp('STORE', exitvar)
+        self.emit('ROT_TWO')
         self.emit('LOAD_ATTR', '__enter__')
         self.emit('CALL_FUNCTION', 0)
         if node.vars is None:
@@ -849,8 +848,6 @@ class CodeGenerator:
         self.emit('LOAD_CONST', None)
         self.nextBlock(final)
         self.setups.push((END_FINALLY, final))
-        self._implicitNameOp('LOAD', exitvar)
-        self._implicitNameOp('DELETE', exitvar)
         self.emit('WITH_CLEANUP')
         self.emit('END_FINALLY')
         self.setups.pop()
diff --git a/Misc/NEWS b/Misc/NEWS
index e3a3cb2..26b8b79 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -12,6 +12,9 @@ What's New in Python 2.6 alpha 2?
 Core and builtins
 -----------------
 
+- Issue #2179: speed up with statement execution by storing the exit method
+  on the stack instead of in a temporary variable (patch by Jeffrey Yaskin)
+
 - Issue #2238: Some syntax errors in *args and **kwargs expressions could give
   bogus error messages.
 
diff --git a/Python/ceval.c b/Python/ceval.c
index f210c9f..4fc1709 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -2254,17 +2254,20 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
 
 		case WITH_CLEANUP:
 		{
-			/* TOP is the context.__exit__ bound method.
-			   Below that are 1-3 values indicating how/why
-			   we entered the finally clause:
-			   - SECOND = None
-			   - (SECOND, THIRD) = (WHY_{RETURN,CONTINUE}), retval
-			   - SECOND = WHY_*; no retval below it
-			   - (SECOND, THIRD, FOURTH) = exc_info()
+			/* At the top of the stack are 1-3 values indicating
+			   how/why we entered the finally clause:
+			   - TOP = None
+			   - (TOP, SECOND) = (WHY_{RETURN,CONTINUE}), retval
+			   - TOP = WHY_*; no retval below it
+			   - (TOP, SECOND, THIRD) = exc_info()
+			   Below them is EXIT, the context.__exit__ bound method.
 			   In the last case, we must call
-			     TOP(SECOND, THIRD, FOURTH)
+			     EXIT(TOP, SECOND, THIRD)
 			   otherwise we must call
-			     TOP(None, None, None)
+			     EXIT(None, None, None)
+
+			   In all cases, we remove EXIT from the stack, leaving
+			   the rest in the same order.
 
 			   In addition, if the stack represents an exception,
 			   *and* the function call returns a 'true' value, we
@@ -2273,36 +2276,59 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
 			   should still be resumed.)
 			*/
 
-			x = TOP();
-			u = SECOND();
-			if (PyInt_Check(u) || u == Py_None) {
+			PyObject *exit_func;
+
+			u = POP();
+			if (u == Py_None) {
+			       	exit_func = TOP();
+				SET_TOP(u);
+				v = w = Py_None;
+			}
+			else if (PyInt_Check(u)) {
+				switch(PyInt_AS_LONG(u)) {
+				case WHY_RETURN:
+				case WHY_CONTINUE:
+					/* Retval in TOP. */
+					exit_func = SECOND();
+					SET_SECOND(TOP());
+					SET_TOP(u);
+					break;
+				default:
+					exit_func = TOP();
+					SET_TOP(u);
+					break;
+				}
 				u = v = w = Py_None;
 			}
 			else {
-				v = THIRD();
-				w = FOURTH();
+				v = TOP();
+				w = SECOND();
+				exit_func = THIRD();
+				SET_TOP(u);
+				SET_SECOND(v);
+				SET_THIRD(w);
 			}
 			/* XXX Not the fastest way to call it... */
-			x = PyObject_CallFunctionObjArgs(x, u, v, w, NULL);
-			if (x == NULL)
+			x = PyObject_CallFunctionObjArgs(exit_func, u, v, w,
+							 NULL);
+			if (x == NULL) {
+				Py_DECREF(exit_func);
 				break; /* Go to error exit */
+			}
 			if (u != Py_None && PyObject_IsTrue(x)) {
 				/* There was an exception and a true return */
-				Py_DECREF(x);
-				x = TOP(); /* Again */
-				STACKADJ(-3);
+				STACKADJ(-2);
 				Py_INCREF(Py_None);
 				SET_TOP(Py_None);
-				Py_DECREF(x);
 				Py_DECREF(u);
 				Py_DECREF(v);
 				Py_DECREF(w);
 			} else {
-				/* Let END_FINALLY do its thing */
-				Py_DECREF(x);
-				x = POP();
-				Py_DECREF(x);
+				/* The stack was rearranged to remove EXIT
+				   above. Let END_FINALLY do its thing */
 			}
+			Py_DECREF(x);
+			Py_DECREF(exit_func);
 			PREDICT(END_FINALLY);
 			break;
 		}
diff --git a/Python/compile.c b/Python/compile.c
index d81fef3..43b7569 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -2842,7 +2842,7 @@ compiler_with(struct compiler *c, stmt_ty s)
 {
     static identifier enter_attr, exit_attr;
     basicblock *block, *finally;
-    identifier tmpexit, tmpvalue = NULL;
+    identifier tmpvalue = NULL;
 
     assert(s->kind == With_kind);
 
@@ -2862,12 +2862,6 @@ compiler_with(struct compiler *c, stmt_ty s)
     if (!block || !finally)
 	return 0;
 
-    /* Create a temporary variable to hold context.__exit__ */
-    tmpexit = compiler_new_tmpname(c);
-    if (tmpexit == NULL)
-	return 0;
-    PyArena_AddPyObject(c->c_arena, tmpexit);
-
     if (s->v.With.optional_vars) {
 	/* Create a temporary variable to hold context.__enter__().
 	   We need to do this rather than preserving it on the stack
@@ -2887,11 +2881,10 @@ compiler_with(struct compiler *c, stmt_ty s)
     /* Evaluate EXPR */
     VISIT(c, expr, s->v.With.context_expr);
 
-    /* Squirrel away context.__exit__  */
+    /* Squirrel away context.__exit__ by stuffing it under context */
     ADDOP(c, DUP_TOP);
     ADDOP_O(c, LOAD_ATTR, exit_attr, names);
-    if (!compiler_nameop(c, tmpexit, Store))
-	return 0;
+    ADDOP(c, ROT_TWO);
 
     /* Call context.__enter__() */
     ADDOP_O(c, LOAD_ATTR, enter_attr, names);
@@ -2935,10 +2928,9 @@ compiler_with(struct compiler *c, stmt_ty s)
     if (!compiler_push_fblock(c, FINALLY_END, finally))
 	return 0;
 
-    /* Finally block starts; push tmpexit and issue our magic opcode. */
-    if (!compiler_nameop(c, tmpexit, Load) ||
-	!compiler_nameop(c, tmpexit, Del))
-	return 0;
+    /* Finally block starts; context.__exit__ is on the stack under
+       the exception or return information. Just issue our magic
+       opcode. */
     ADDOP(c, WITH_CLEANUP);
 
     /* Finally block ends. */
diff --git a/Python/import.c b/Python/import.c
index 191c039..ecbec15 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -72,9 +72,10 @@ extern time_t PyOS_GetLastModificationTime(char *, FILE *);
        			    storing constants that should have been removed)
        Python 2.5c2: 62131 (fix wrong code: for x, in ... in listcomp/genexp)
        Python 2.6a0: 62151 (peephole optimizations and STORE_MAP opcode)
+       Python 2.6a1: 62161 (WITH_CLEANUP optimization)
 .
 */
-#define MAGIC (62151 | ((long)'\r'<<16) | ((long)'\n'<<24))
+#define MAGIC (62161 | ((long)'\r'<<16) | ((long)'\n'<<24))
 
 /* Magic word as global; note that _PyImport_Init() can change the
    value of this global to accommodate for alterations of how the