py: Modify bytecode "with" behaviour so it doesn't use any heap.

dpgeorge · dpgeorge · commit 8c1d23a0e20d · 2015-04-24T01:52:28.000+01:00
Before this patch a "with" block needed to create a bound method object
on the heap for the __exit__ call.  Now it doesn't because we use
load_method instead of load_attr, and save the method+self on the stack.
diff --git a/py/emitbc.c b/py/emitbc.c
@@ -680,12 +680,15 @@ void mp_emit_bc_unwind_jump(emit_t *emit, mp_uint_t label, mp_uint_t except_dept
 }
 
 void mp_emit_bc_setup_with(emit_t *emit, mp_uint_t label) {
-    emit_bc_pre(emit, 7);
+    // TODO We can probably optimise the amount of needed stack space, since
+    // we don't actually need 4 slots during the entire with block, only in
+    // the cleanup handler in certain cases.  It needs some thinking.
+    emit_bc_pre(emit, 4);
     emit_write_bytecode_byte_unsigned_label(emit, MP_BC_SETUP_WITH, label);
 }
 
 void mp_emit_bc_with_cleanup(emit_t *emit) {
-    emit_bc_pre(emit, -7);
+    emit_bc_pre(emit, -4);
     emit_write_bytecode_byte(emit, MP_BC_WITH_CLEANUP);
 }
 
diff --git a/py/vm.c b/py/vm.c
@@ -548,68 +548,82 @@ run_code_state: ;
 
                 ENTRY(MP_BC_SETUP_WITH): {
                     MARK_EXC_IP_SELECTIVE();
+                    // stack: (..., ctx_mgr)
                     mp_obj_t obj = TOP();
-                    SET_TOP(mp_load_attr(obj, MP_QSTR___exit__));
-                    mp_load_method(obj, MP_QSTR___enter__, sp + 1);
-                    mp_obj_t ret = mp_call_method_n_kw(0, 0, sp + 1);
+                    mp_load_method(obj, MP_QSTR___exit__, sp);
+                    mp_load_method(obj, MP_QSTR___enter__, sp + 2);
+                    mp_obj_t ret = mp_call_method_n_kw(0, 0, sp + 2);
+                    sp += 1;
                     PUSH_EXC_BLOCK(1);
                     PUSH(ret);
+                    // stack: (..., __exit__, ctx_mgr, as_value)
                     DISPATCH();
                 }
 
                 ENTRY(MP_BC_WITH_CLEANUP): {
                     MARK_EXC_IP_SELECTIVE();
                     // Arriving here, there's "exception control block" on top of stack,
-                    // and __exit__ bound method underneath it. Bytecode calls __exit__,
+                    // and __exit__ method (with self) underneath it. Bytecode calls __exit__,
                     // and "deletes" it off stack, shifting "exception control block"
                     // to its place.
-                    static const mp_obj_t no_exc[] = {mp_const_none, mp_const_none, mp_const_none};
                     if (TOP() == mp_const_none) {
-                        sp--;
-                        mp_obj_t obj = TOP();
+                        // stack: (..., __exit__, ctx_mgr, None)
+                        sp[1] = mp_const_none;
+                        sp[2] = mp_const_none;
+                        sp -= 2;
+                        mp_call_method_n_kw(3, 0, sp);
                         SET_TOP(mp_const_none);
-                        mp_call_function_n_kw(obj, 3, 0, no_exc);
                     } else if (MP_OBJ_IS_SMALL_INT(TOP())) {
                         mp_int_t cause_val = MP_OBJ_SMALL_INT_VALUE(TOP());
                         if (cause_val == UNWIND_RETURN) {
-                            mp_call_function_n_kw(sp[-2], 3, 0, no_exc);
+                            // stack: (..., __exit__, ctx_mgr, ret_val, UNWIND_RETURN)
+                            mp_obj_t ret_val = sp[-1];
+                            sp[-1] = mp_const_none;
+                            sp[0] = mp_const_none;
+                            sp[1] = mp_const_none;
+                            mp_call_method_n_kw(3, 0, sp - 3);
+                            sp[-3] = ret_val;
+                            sp[-2] = MP_OBJ_NEW_SMALL_INT(UNWIND_RETURN);
                         } else {
                             assert(cause_val == UNWIND_JUMP);
-                            mp_call_function_n_kw(sp[-3], 3, 0, no_exc);
-                            // Pop __exit__ boundmethod at sp[-3]
-                            sp[-3] = sp[-2];
+                            // stack: (..., __exit__, ctx_mgr, dest_ip, num_exc, UNWIND_JUMP)
+                            mp_obj_t dest_ip = sp[-2];
+                            mp_obj_t num_exc = sp[-1];
+                            sp[-2] = mp_const_none;
+                            sp[-1] = mp_const_none;
+                            sp[0] = mp_const_none;
+                            mp_call_method_n_kw(3, 0, sp - 4);
+                            sp[-4] = dest_ip;
+                            sp[-3] = num_exc;
+                            sp[-2] = MP_OBJ_NEW_SMALL_INT(UNWIND_JUMP);
                         }
-                        sp[-2] = sp[-1]; // copy retval down
-                        sp[-1] = sp[0]; // copy cause down
-                        sp--; // discard top value (was cause)
+                        sp -= 2; // we removed (__exit__, ctx_mgr)
                     } else {
                         assert(mp_obj_is_exception_type(TOP()));
+                        // stack: (..., __exit__, ctx_mgr, traceback, exc_val, exc_type)
                         // Need to pass (sp[0], sp[-1], sp[-2]) as arguments so must reverse the
                         // order of these on the value stack (don't want to create a temporary
                         // array because it increases stack footprint of the VM).
                         mp_obj_t obj = sp[-2];
                         sp[-2] = sp[0];
                         sp[0] = obj;
-                        mp_obj_t ret_value = mp_call_function_n_kw(sp[-3], 3, 0, &sp[-2]);
+                        mp_obj_t ret_value = mp_call_method_n_kw(3, 0, sp - 4);
                         if (mp_obj_is_true(ret_value)) {
-                            // 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...
+                            // We need to silence/swallow the exception.  This is done
+                            // by popping the exception and the __exit__ handler and
+                            // replacing it with None, which signals END_FINALLY to just
+                            // execute the finally handler normally.
                             sp -= 4;
-                            // ... pop "with" exception handler, and signal END_FINALLY
-                            // to just execute finally handler normally (by pushing None
-                            // on value stack)
+                            SET_TOP(mp_const_none);
                             assert(exc_sp >= exc_stack);
                             POP_EXC_BLOCK();
-                            PUSH(mp_const_none);
                         } else {
-                            // Pop __exit__ boundmethod at sp[-3], remembering that top 3 values
-                            // are reversed.
-                            sp[-3] = sp[0];
-                            obj = sp[-2];
-                            sp[-2] = sp[-1];
-                            sp[-1] = obj;
-                            sp--;
+                            // We need to re-raise the exception.  We pop __exit__ handler
+                            // and copy the 3 exception values down (remembering that they
+                            // are reversed due to above code).
+                            sp[-4] = sp[0];
+                            sp[-3] = sp[-1];
+                            sp -= 2;
                         }
                     }
                     DISPATCH();
diff --git a/tests/cmdline/cmd_showbc.py.exp b/tests/cmdline/cmd_showbc.py.exp
@@ -29,7 +29,7 @@ Raw bytecode (code_info_size=\\d\+, bytecode_size=\\d\+):
 ########
 \.\+5b
 arg names:
-(N_STATE 25)
+(N_STATE 22)
 (N_EXC_STACK 2)
 (INIT_CELL 14)
 (INIT_CELL 15)

Original file line number	Diff line number	Diff line change
`@@ -680,12 +680,15 @@ void mp_emit_bc_unwind_jump(emit_t *emit, mp_uint_t label, mp_uint_t except_dept`
`680`	`680`	`}`
`681`	`681`
`682`	`682`	`void mp_emit_bc_setup_with(emit_t *emit, mp_uint_t label) {`
`683`		`- emit_bc_pre(emit, 7);`
	`683`	`+ // TODO We can probably optimise the amount of needed stack space, since`
	`684`	`+ // we don't actually need 4 slots during the entire with block, only in`
	`685`	`+ // the cleanup handler in certain cases. It needs some thinking.`
	`686`	`+ emit_bc_pre(emit, 4);`
`684`	`687`	`emit_write_bytecode_byte_unsigned_label(emit, MP_BC_SETUP_WITH, label);`
`685`	`688`	`}`
`686`	`689`
`687`	`690`	`void mp_emit_bc_with_cleanup(emit_t *emit) {`
`688`		`- emit_bc_pre(emit, -7);`
	`691`	`+ emit_bc_pre(emit, -4);`
`689`	`692`	`emit_write_bytecode_byte(emit, MP_BC_WITH_CLEANUP);`
`690`	`693`	`}`
`691`	`694`