From 4bf3f2d3c0d7769b5a771e58a06c13ee527e657d Mon Sep 17 00:00:00 2001 From: Damien George Date: Thu, 15 Oct 2015 17:48:28 +0100 Subject: [PATCH] py: Fix with+for+return bug by popping for-iter when unwinding exc stack. Addresses issue #1182. --- py/vm.c | 23 +++++++++++--- tests/basics/try_finally_return.py | 49 ++++++++++++++++++++++++++++++ tests/basics/with_return.py | 47 +++++++++++++++++++++++++--- 3 files changed, 110 insertions(+), 9 deletions(-) diff --git a/py/vm.c b/py/vm.c index 2611be683..393b8a1db 100644 --- a/py/vm.c +++ b/py/vm.c @@ -637,10 +637,14 @@ unwind_jump:; unum -= 1; assert(exc_sp >= exc_stack); if (MP_TAGPTR_TAG1(exc_sp->val_sp)) { + // Getting here the stack looks like: + // (..., X, dest_ip) + // where X is pointed to by exc_sp->val_sp and in the case + // of a "with" block contains the context manager info. // We're going to run "finally" code as a coroutine // (not calling it recursively). Set up a sentinel // on a stack so it can return back to us when it is - // done (when END_FINALLY reached). + // done (when WITH_CLEANUP or END_FINALLY reached). PUSH((void*)unum); // push number of exception handlers left to unwind PUSH(MP_OBJ_NEW_SMALL_INT(UNWIND_JUMP)); // push sentinel ip = exc_sp->handler; // get exception handler byte code address @@ -1016,15 +1020,24 @@ unwind_jump:; unwind_return: while (exc_sp >= exc_stack) { if (MP_TAGPTR_TAG1(exc_sp->val_sp)) { + // Getting here the stack looks like: + // (..., X, [iter0, iter1, ...,] ret_val) + // where X is pointed to by exc_sp->val_sp and in the case + // of a "with" block contains the context manager info. + // There may be 0 or more for-iterators between X and the + // return value, and these must be removed before control can + // pass to the finally code. We simply copy the ret_value down + // over these iterators, if they exist. If they don't then the + // following is a null operation. + mp_obj_t *finally_sp = MP_TAGPTR_PTR(exc_sp->val_sp); + finally_sp[1] = sp[0]; + sp = &finally_sp[1]; // We're going to run "finally" code as a coroutine // (not calling it recursively). Set up a sentinel // on a stack so it can return back to us when it is - // done (when END_FINALLY reached). + // done (when WITH_CLEANUP or END_FINALLY reached). PUSH(MP_OBJ_NEW_SMALL_INT(UNWIND_RETURN)); ip = exc_sp->handler; - // We don't need to do anything with sp, finally is just - // syntactic sugar for sequential execution?? - // sp = exc_sp--; goto dispatch_loop; } diff --git a/tests/basics/try_finally_return.py b/tests/basics/try_finally_return.py index 4adf3f097..31a507e8d 100644 --- a/tests/basics/try_finally_return.py +++ b/tests/basics/try_finally_return.py @@ -21,3 +21,52 @@ def func3(): print("finally 3") print(func3()) + +# for loop within try-finally +def f(): + try: + for i in [1, 2]: + return i + finally: + print('finally') +print(f()) + +# multiple for loops within try-finally +def f(): + try: + for i in [1, 2]: + for j in [3, 4]: + return (i, j) + finally: + print('finally') +print(f()) + +# multiple for loops and nested try-finally's +def f(): + try: + for i in [1, 2]: + for j in [3, 4]: + try: + for k in [5, 6]: + for l in [7, 8]: + return (i, j, k, l) + finally: + print('finally 2') + finally: + print('finally 1') +print(f()) + +# multiple for loops that are optimised, and nested try-finally's +def f(): + try: + for i in range(1, 3): + for j in range(3, 5): + try: + for k in range(5, 7): + for l in range(7, 9): + return (i, j, k, l) + finally: + print('finally 2') + finally: + print('finally 1') +print(f()) diff --git a/tests/basics/with_return.py b/tests/basics/with_return.py index cb0135c8b..fd848f133 100644 --- a/tests/basics/with_return.py +++ b/tests/basics/with_return.py @@ -1,14 +1,53 @@ class CtxMgr: + def __init__(self, id): + self.id = id def __enter__(self): - print("__enter__") + print("__enter__", self.id) return self def __exit__(self, a, b, c): - print("__exit__", repr(a), repr(b)) + print("__exit__", self.id, repr(a), repr(b)) +# simple case def foo(): - with CtxMgr(): + with CtxMgr(1): return 4 - print(foo()) + +# for loop within with (iterator needs removing upon return) +def f(): + with CtxMgr(1): + for i in [1, 2]: + return i +print(f()) + +# multiple for loops within with +def f(): + with CtxMgr(1): + for i in [1, 2]: + for j in [3, 4]: + return (i, j) +print(f()) + +# multiple for loops within nested withs +def f(): + with CtxMgr(1): + for i in [1, 2]: + for j in [3, 4]: + with CtxMgr(2): + for k in [5, 6]: + for l in [7, 8]: + return (i, j, k, l) +print(f()) + +# multiple for loops that are optimised, and nested withs +def f(): + with CtxMgr(1): + for i in range(1, 3): + for j in range(3, 5): + with CtxMgr(2): + for k in range(5, 7): + for l in range(7, 9): + return (i, j, k, l) +print(f())