py/emitnative: Fix native async with.

The code generating the entry to the finally handler of an async-with
statement was simply wrong for the case of the native emitter.  Among other
things the layout of the stack was incorrect.

This is fixed by this commit.  The setup of the async-with finally handler
is now put in a dedicated emit function, for both the bytecode and native
emitters to implement in their own way (the bytecode emitter is unchanged,
just factored to a function).

With this fix all of the async-with tests now work when using the native
emitter.

Signed-off-by: Damien George <damien@micropython.org>

1 files changed, 5 insertions, 13 deletions

diff --git a/py/compile.c b/py/compile.c
index 3429e384c..d2af6aaf2 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -1899,19 +1899,7 @@ static void compile_async_with_stmt_helper(compiler_t *comp, size_t n, mp_parse_
 
         // Handle case 1: call __aexit__
         // Stack: (..., ctx_mgr)
-        EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE); // to tell end_finally there's no exception
-        EMIT(rot_two);
-        EMIT_ARG(jump, l_aexit_no_exc); // jump to code below to call __aexit__
-
-        // Start of "finally" block
-        // At this point we have case 2 or 3, we detect which one by the TOS being an exception or not
-        EMIT_ARG(label_assign, l_finally_block);
-
-        // Detect if TOS an exception or not
-        EMIT(dup_top);
-        EMIT_LOAD_GLOBAL(MP_QSTR_BaseException);
-        EMIT_ARG(binary_op, MP_BINARY_OP_EXCEPTION_MATCH);
-        EMIT_ARG(pop_jump_if, false, l_ret_unwind_jump); // if not an exception then we have case 3
+        EMIT_ARG(async_with_setup_finally, l_aexit_no_exc, l_finally_block, l_ret_unwind_jump);
 
         // Handle case 2: call __aexit__ and either swallow or re-raise the exception
         // Stack: (..., ctx_mgr, exc)
@@ -1937,6 +1925,7 @@ static void compile_async_with_stmt_helper(compiler_t *comp, size_t n, mp_parse_
         EMIT_ARG(pop_jump_if, false, l_end);
         EMIT(pop_top); // pop exception
         EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE); // replace with None to swallow exception
+        // Stack: (..., None)
         EMIT_ARG(jump, l_end);
         EMIT_ARG(adjust_stack_size, 2);
 
@@ -1946,6 +1935,8 @@ static void compile_async_with_stmt_helper(compiler_t *comp, size_t n, mp_parse_
         EMIT(rot_three);
         EMIT(rot_three);
         EMIT_ARG(label_assign, l_aexit_no_exc);
+        // We arrive here from either case 1 (a jump) or case 3 (fall through)
+        // Stack: case 1: (..., None, ctx_mgr) or case 3: (..., X, INT, ctx_mgr)
         EMIT_ARG(load_method, MP_QSTR___aexit__, false);
         EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
         EMIT(dup_top);
@@ -1953,6 +1944,7 @@ static void compile_async_with_stmt_helper(compiler_t *comp, size_t n, mp_parse_
         EMIT_ARG(call_method, 3, 0, 0);
         compile_yield_from(comp);
         EMIT(pop_top);
+        // Stack: case 1: (..., None) or case 3: (..., X, INT)
         EMIT_ARG(adjust_stack_size, -1);
 
         // End of "finally" block