py: Implement basic with support in native emitter.

dpgeorge · dpgeorge · commit 2c915e1ae617 · 2016-04-07T08:53:24.000+01:00
diff --git a/py/compile.c b/py/compile.c
@@ -1624,6 +1624,11 @@ STATIC void compile_with_stmt_helper(compiler_t *comp, int n, mp_parse_node_t *n
         compile_node(comp, body);
     } else {
         uint l_end = comp_next_label(comp);
+        if (MICROPY_EMIT_NATIVE && comp->scope_cur->emit_options != MP_EMIT_OPT_BYTECODE) {
+            // we need to allocate an extra label for the native emitter
+            // it will use l_end+1 as an auxiliary label
+            comp_next_label(comp);
+        }
         if (MP_PARSE_NODE_IS_STRUCT_KIND(nodes[0], PN_with_item)) {
             // this pre-bit is of the form "a as b"
             mp_parse_node_struct_t *pns = (mp_parse_node_struct_t*)nodes[0];
diff --git a/py/emitnative.c b/py/emitnative.c
@@ -1984,15 +1984,114 @@ STATIC void emit_native_continue_loop(emit_t *emit, mp_uint_t label, mp_uint_t e
 }
 
 STATIC void emit_native_setup_with(emit_t *emit, mp_uint_t label) {
-    // not supported, or could be with runtime call
-    (void)emit;
-    (void)label;
-    assert(0);
+    // the context manager is on the top of the stack
+    // stack: (..., ctx_mgr)
+
+    // get __exit__ method
+    vtype_kind_t vtype;
+    emit_access_stack(emit, 1, &vtype, REG_ARG_1); // arg1 = ctx_mgr
+    assert(vtype == VTYPE_PYOBJ);
+    emit_get_stack_pointer_to_reg_for_push(emit, REG_ARG_3, 2); // arg3 = dest ptr
+    emit_call_with_imm_arg(emit, MP_F_LOAD_METHOD, MP_QSTR___exit__, REG_ARG_2);
+    // stack: (..., ctx_mgr, __exit__, self)
+
+    emit_pre_pop_reg(emit, &vtype, REG_ARG_3); // self
+    emit_pre_pop_reg(emit, &vtype, REG_ARG_2); // __exit__
+    emit_pre_pop_reg(emit, &vtype, REG_ARG_1); // ctx_mgr
+    emit_post_push_reg(emit, vtype, REG_ARG_2); // __exit__
+    emit_post_push_reg(emit, vtype, REG_ARG_3); // self
+    // stack: (..., __exit__, self)
+    // REG_ARG_1=ctx_mgr
+
+    // get __enter__ method
+    emit_get_stack_pointer_to_reg_for_push(emit, REG_ARG_3, 2); // arg3 = dest ptr
+    emit_call_with_imm_arg(emit, MP_F_LOAD_METHOD, MP_QSTR___enter__, REG_ARG_2); // arg2 = method name
+    // stack: (..., __exit__, self, __enter__, self)
+
+    // call __enter__ method
+    emit_get_stack_pointer_to_reg_for_pop(emit, REG_ARG_3, 2); // pointer to items, including meth and self
+    emit_call_with_2_imm_args(emit, MP_F_CALL_METHOD_N_KW, 0, REG_ARG_1, 0, REG_ARG_2);
+    emit_post_push_reg(emit, VTYPE_PYOBJ, REG_RET); // push return value of __enter__
+    // stack: (..., __exit__, self, as_value)
+
+    // need to commit stack because we may jump elsewhere
+    need_stack_settled(emit);
+    emit_get_stack_pointer_to_reg_for_push(emit, REG_ARG_1, sizeof(nlr_buf_t) / sizeof(mp_uint_t)); // arg1 = pointer to nlr buf
+    emit_call(emit, MP_F_NLR_PUSH);
+    ASM_JUMP_IF_REG_NONZERO(emit->as, REG_RET, label);
+
+    emit_access_stack(emit, sizeof(nlr_buf_t) / sizeof(mp_uint_t) + 1, &vtype, REG_RET); // access return value of __enter__
+    emit_post_push_reg(emit, VTYPE_PYOBJ, REG_RET); // push return value of __enter__
+    // stack: (..., __exit__, self, as_value, nlr_buf, as_value)
 }
 
 STATIC void emit_native_with_cleanup(emit_t *emit, mp_uint_t label) {
-    (void)emit;
-    assert(0);
+    // note: label+1 is available as an auxiliary label
+
+    // stack: (..., __exit__, self, as_value, nlr_buf)
+    emit_native_pre(emit);
+    emit_call(emit, MP_F_NLR_POP);
+    adjust_stack(emit, -(mp_int_t)(sizeof(nlr_buf_t) / sizeof(mp_uint_t)) - 1);
+    // stack: (..., __exit__, self)
+
+    // call __exit__
+    emit_post_push_imm(emit, VTYPE_PYOBJ, (mp_uint_t)mp_const_none);
+    emit_post_push_imm(emit, VTYPE_PYOBJ, (mp_uint_t)mp_const_none);
+    emit_post_push_imm(emit, VTYPE_PYOBJ, (mp_uint_t)mp_const_none);
+    emit_get_stack_pointer_to_reg_for_pop(emit, REG_ARG_3, 5);
+    emit_call_with_2_imm_args(emit, MP_F_CALL_METHOD_N_KW, 3, REG_ARG_1, 0, REG_ARG_2);
+
+    // jump to after with cleanup nlr_catch block
+    adjust_stack(emit, 1); // dummy nlr_buf.prev
+    emit_native_load_const_tok(emit, MP_TOKEN_KW_NONE); // nlr_buf.ret_val = no exception
+    emit_native_jump(emit, label + 1);
+
+    // nlr_catch
+    emit_native_label_assign(emit, label);
+
+    // adjust stack counter for: __exit__, self, as_value
+    adjust_stack(emit, 3);
+    // stack: (..., __exit__, self, as_value, nlr_buf.prev, nlr_buf.ret_val)
+
+    vtype_kind_t vtype;
+    emit_pre_pop_reg(emit, &vtype, REG_ARG_1); // get the thrown value (exc)
+    adjust_stack(emit, -2); // discard nlr_buf.prev and as_value
+    // stack: (..., __exit__, self)
+    // REG_ARG_1=exc
+
+    emit_pre_pop_reg(emit, &vtype, REG_ARG_2); // self
+    emit_pre_pop_reg(emit, &vtype, REG_ARG_3); // __exit__
+    adjust_stack(emit, 1); // dummy nlr_buf.prev
+    emit_post_push_reg(emit, vtype, REG_ARG_1); // push exc to save it for later
+    emit_post_push_reg(emit, vtype, REG_ARG_3); // __exit__
+    emit_post_push_reg(emit, vtype, REG_ARG_2); // self
+    // stack: (..., exc, __exit__, self)
+    // REG_ARG_1=exc
+
+    ASM_LOAD_REG_REG_OFFSET(emit->as, REG_ARG_2, REG_ARG_1, 0); // get type(exc)
+    emit_post_push_reg(emit, VTYPE_PYOBJ, REG_ARG_2); // push type(exc)
+    emit_post_push_reg(emit, VTYPE_PYOBJ, REG_ARG_1); // push exc value
+    emit_post_push_imm(emit, VTYPE_PYOBJ, (mp_uint_t)mp_const_none); // traceback info
+    // stack: (..., exc, __exit__, self, type(exc), exc, traceback)
+
+    // call __exit__ method
+    emit_get_stack_pointer_to_reg_for_pop(emit, REG_ARG_3, 5);
+    emit_call_with_2_imm_args(emit, MP_F_CALL_METHOD_N_KW, 3, REG_ARG_1, 0, REG_ARG_2);
+    // stack: (..., exc)
+
+    // if REG_RET is true then we need to replace top-of-stack with None (swallow exception)
+    if (REG_ARG_1 != REG_RET) {
+        ASM_MOV_REG_REG(emit->as, REG_ARG_1, REG_RET);
+    }
+    emit_call(emit, MP_F_OBJ_IS_TRUE);
+    ASM_JUMP_IF_REG_ZERO(emit->as, REG_RET, label + 1);
+
+    // replace exc with None
+    emit_pre_pop_discard(emit);
+    emit_post_push_imm(emit, VTYPE_PYOBJ, (mp_uint_t)mp_const_none);
+
+    // end of with cleanup nlr_catch block
+    emit_native_label_assign(emit, label + 1);
 }
 
 STATIC void emit_native_setup_except(emit_t *emit, mp_uint_t label) {
@@ -2016,7 +2115,8 @@ STATIC void emit_native_end_finally(emit_t *emit) {
     //   else: raise exc
     // the check if exc is None is done in the MP_F_NATIVE_RAISE stub
     vtype_kind_t vtype;
-    emit_pre_pop_reg(emit, &vtype, REG_ARG_1);
+    emit_pre_pop_reg(emit, &vtype, REG_ARG_1); // get nlr_buf.ret_val
+    emit_pre_pop_discard(emit); // discard nlr_buf.prev
     emit_call(emit, MP_F_NATIVE_RAISE);
     emit_post(emit);
 }
@@ -2053,7 +2153,7 @@ STATIC void emit_native_for_iter_end(emit_t *emit) {
 STATIC void emit_native_pop_block(emit_t *emit) {
     emit_native_pre(emit);
     emit_call(emit, MP_F_NLR_POP);
-    adjust_stack(emit, -(mp_int_t)(sizeof(nlr_buf_t) / sizeof(mp_uint_t)));
+    adjust_stack(emit, -(mp_int_t)(sizeof(nlr_buf_t) / sizeof(mp_uint_t)) + 1);
     emit_post(emit);
 }
 
@@ -2495,15 +2595,15 @@ STATIC void emit_native_start_except_handler(emit_t *emit) {
     // This instruction follows an nlr_pop, so the stack counter is back to zero, when really
     // it should be up by a whole nlr_buf_t.  We then want to pop the nlr_buf_t here, but save
     // the first 2 elements, so we can get the thrown value.
-    adjust_stack(emit, 2);
+    adjust_stack(emit, 1);
     vtype_kind_t vtype_nlr;
     emit_pre_pop_reg(emit, &vtype_nlr, REG_ARG_1); // get the thrown value
     emit_pre_pop_discard(emit); // discard the linked-list pointer in the nlr_buf
     emit_post_push_reg_reg_reg(emit, VTYPE_PYOBJ, REG_ARG_1, VTYPE_PYOBJ, REG_ARG_1, VTYPE_PYOBJ, REG_ARG_1); // push the 3 exception items
 }
 
 STATIC void emit_native_end_except_handler(emit_t *emit) {
-    adjust_stack(emit, -2);
+    adjust_stack(emit, -1);
 }
 
 const emit_method_table_t EXPORT_FUN(method_table) = {
diff --git a/tests/run-tests b/tests/run-tests
@@ -244,7 +244,7 @@ def run_tests(pyb, tests, args):
         skip_tests.update({'basics/%s.py' % t for t in 'gen_yield_from gen_yield_from_close gen_yield_from_ducktype gen_yield_from_exc gen_yield_from_iter gen_yield_from_send gen_yield_from_throw generator1 generator2 generator_args generator_close generator_closure generator_exc generator_return generator_send'.split()}) # require yield
         skip_tests.update({'basics/%s.py' % t for t in 'bytes_gen class_store_class globals_del string_join'.split()}) # require yield
         skip_tests.update({'basics/%s.py' % t for t in 'try_reraise try_reraise2'.split()}) # require raise_varargs
-        skip_tests.update({'basics/%s.py' % t for t in 'with1 with_break with_continue with_return'.split()}) # require with
+        skip_tests.update({'basics/%s.py' % t for t in 'with_break with_continue with_return'.split()}) # require complete with support
         skip_tests.add('basics/array_construct2.py') # requires generators
         skip_tests.add('basics/bool1.py') # seems to randomly fail
         skip_tests.add('basics/class_bind_self.py') # requires yield
@@ -257,8 +257,6 @@ def run_tests(pyb, tests, args):
         skip_tests.add('basics/try_finally_return2.py') # requires proper try finally code
         skip_tests.add('basics/unboundlocal.py') # requires checking for unbound local
         skip_tests.add('import/gen_context.py') # requires yield_value
-        skip_tests.add('io/file_with.py') # requires with
-        skip_tests.add('io/stringio_with.py') # requires with
         skip_tests.add('misc/features.py') # requires raise_varargs
         skip_tests.add('misc/rge_sm.py') # requires yield
         skip_tests.add('misc/print_exception.py') # because native doesn't have proper traceback info
