vm: Implement "with" statement (SETUP_WITH and WITH_CLEANUP bytecodes).

3 files changed, 73 insertions, 2 deletions

diff --git a/py/compile.c b/py/compile.c
index b3a83715e0..9e6c4e5fe4 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -1798,6 +1798,7 @@ void compile_with_stmt_helper(compiler_t *comp, int n, mp_parse_node_t *nodes, m
             EMIT_ARG(setup_with, l_end);
             EMIT(pop_top);
         }
+        compile_increase_except_level(comp);
         // compile additional pre-bits and the body
         compile_with_stmt_helper(comp, n - 1, nodes + 1, body);
         // finish this with block
@@ -1805,6 +1806,7 @@ void compile_with_stmt_helper(compiler_t *comp, int n, mp_parse_node_t *nodes, m
         EMIT_ARG(load_const_tok, MP_TOKEN_KW_NONE);
         EMIT_ARG(label_assign, l_end);
         EMIT(with_cleanup);
+        compile_decrease_except_level(comp);
         EMIT(end_finally);
     }
 }
diff --git a/py/qstrdefs.h b/py/qstrdefs.h
index 3be6295067..457043938b 100644
--- a/py/qstrdefs.h
+++ b/py/qstrdefs.h
@@ -16,6 +16,8 @@ Q(__qualname__)
 Q(__repl_print__)
 
 Q(__bool__)
+Q(__enter__)
+Q(__exit__)
 Q(__len__)
 Q(__iter__)
 Q(__getitem__)
diff --git a/py/vm.c b/py/vm.c
index 60ed641a85..206e9c8ce3 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -378,6 +378,73 @@ dispatch_loop:
                         break;
                         */
 
+                    case MP_BC_SETUP_WITH: {
+                        obj1 = TOP();
+                        SET_TOP(rt_load_attr(obj1, MP_QSTR___exit__));
+                        mp_obj_t dest[2];
+                        rt_load_method(obj1, MP_QSTR___enter__, dest);
+                        obj2 = rt_call_method_n_kw(0, 0, dest);
+                        SETUP_BLOCK();
+                        PUSH(obj2);
+                        break;
+                    }
+
+                    case MP_BC_WITH_CLEANUP: {
+                        static const mp_obj_t no_exc[] = {mp_const_none, mp_const_none, mp_const_none};
+                        if (TOP() == mp_const_none) {
+                            sp--;
+                            obj1 = TOP();
+                            SET_TOP(mp_const_none);
+                            obj2 = rt_call_function_n_kw(obj1, 3, 0, no_exc);
+                        } else if (MP_OBJ_IS_SMALL_INT(TOP())) {
+                            mp_obj_t cause = POP();
+                            switch (MP_OBJ_SMALL_INT_VALUE(cause)) {
+                                case UNWIND_RETURN: {
+                                    mp_obj_t retval = POP();
+                                    obj2 = rt_call_function_n_kw(TOP(), 3, 0, no_exc);
+                                    SET_TOP(retval);
+                                    PUSH(cause);
+                                    break;
+                                }
+                                case UNWIND_JUMP: {
+                                    obj2 = rt_call_function_n_kw(sp[-2], 3, 0, no_exc);
+                                    // Pop __exit__ boundmethod at sp[-2]
+                                    sp[-2] = sp[-1];
+                                    sp[-1] = sp[0];
+                                    SET_TOP(cause);
+                                    break;
+                                }
+                                default:
+                                    assert(0);
+                            }
+                        } else if (mp_obj_is_exception_type(TOP())) {
+                            mp_obj_t args[3] = {sp[0], sp[-1], sp[-2]};
+                            obj2 = rt_call_function_n_kw(sp[-3], 3, 0, args);
+                            // Pop __exit__ boundmethod at sp[-3]
+                            // TODO: Once semantics is proven, optimize for case when obj2 == True
+                            sp[-3] = sp[-2];
+                            sp[-2] = sp[-1];
+                            sp[-1] = sp[0];
+                            sp--;
+                            if (rt_is_true(obj2)) {
+                                // This is what CPython does
+                                //PUSH(MP_OBJ_NEW_SMALL_INT(UNWIND_SILENCED));
+                                // But what we need to do is - pop exception from value stack...
+                                sp -= 3;
+                                // ... pop with exception handler, and signal END_FINALLY
+                                // to just execute finally handler normally (signalled by None
+                                // on value stack)
+                                assert(exc_sp >= exc_stack);
+                                assert(exc_sp->opcode == MP_BC_SETUP_WITH);
+                                exc_sp--;
+                                PUSH(mp_const_none);
+                            }
+                        } else {
+                            assert(0);
+                        }
+                        break;
+                    }
+
                     case MP_BC_UNWIND_JUMP:
                         DECODE_SLABEL;
                         PUSH((void*)(ip + unum)); // push destination ip for jump
@@ -387,7 +454,7 @@ unwind_jump:
                         while (unum > 0) {
                             unum -= 1;
                             assert(exc_sp >= exc_stack);
-                            if (exc_sp->opcode == MP_BC_SETUP_FINALLY) {
+                            if (exc_sp->opcode == MP_BC_SETUP_FINALLY || exc_sp->opcode == MP_BC_SETUP_WITH) {
                                 // 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
@@ -601,7 +668,7 @@ unwind_jump:
                     case MP_BC_RETURN_VALUE:
 unwind_return:
                         while (exc_sp >= exc_stack) {
-                            if (exc_sp->opcode == MP_BC_SETUP_FINALLY) {
+                            if (exc_sp->opcode == MP_BC_SETUP_FINALLY || exc_sp->opcode == MP_BC_SETUP_WITH) {
                                 // 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