py: Use nlr jump callbacks to optimise compile/execute functions.

The changed functions now use less stack, and don't have any issues with
local variables needing to be declared volatile.

Testing on a PYBv1.0, imports (of .py, .mpy and frozen code) now use 64
less bytes of C stack per import depth.

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

1 files changed, 16 insertions, 16 deletions

diff --git a/py/builtinevex.c b/py/builtinevex.c
index 173978ef5..97eab7fad 100644
--- a/py/builtinevex.c
+++ b/py/builtinevex.c
@@ -45,12 +45,18 @@ STATIC MP_DEFINE_CONST_OBJ_TYPE(
     );
 
 STATIC mp_obj_t code_execute(mp_obj_code_t *self, mp_obj_dict_t *globals, mp_obj_dict_t *locals) {
-    // save context and set new context
-    mp_obj_dict_t *old_globals = mp_globals_get();
-    mp_obj_dict_t *old_locals = mp_locals_get();
+    // save context
+    nlr_jump_callback_node_globals_locals_t ctx;
+    ctx.globals = mp_globals_get();
+    ctx.locals = mp_locals_get();
+
+    // set new context
     mp_globals_set(globals);
     mp_locals_set(locals);
 
+    // set exception handler to restore context if an exception is raised
+    nlr_push_jump_callback(&ctx.callback, mp_globals_locals_set_from_nlr_jump_callback);
+
     // a bit of a hack: fun_bc will re-set globals, so need to make sure it's
     // the correct one
     if (mp_obj_is_type(self->module_fun, &mp_type_fun_bc)) {
@@ -59,19 +65,13 @@ STATIC mp_obj_t code_execute(mp_obj_code_t *self, mp_obj_dict_t *globals, mp_obj
     }
 
     // execute code
-    nlr_buf_t nlr;
-    if (nlr_push(&nlr) == 0) {
-        mp_obj_t ret = mp_call_function_0(self->module_fun);
-        nlr_pop();
-        mp_globals_set(old_globals);
-        mp_locals_set(old_locals);
-        return ret;
-    } else {
-        // exception; restore context and re-raise same exception
-        mp_globals_set(old_globals);
-        mp_locals_set(old_locals);
-        nlr_jump(nlr.ret_val);
-    }
+    mp_obj_t ret = mp_call_function_0(self->module_fun);
+
+    // deregister exception handler and restore context
+    nlr_pop_jump_callback(true);
+
+    // return value
+    return ret;
 }
 
 STATIC mp_obj_t mp_builtin_compile(size_t n_args, const mp_obj_t *args) {