From 3504edc8048f8ab038ace50b0bbdf65b30619cc5 Mon Sep 17 00:00:00 2001
From: Damien George <damien.p.george@gmail.com>
Date: Tue, 10 Sep 2019 13:44:22 +1000
Subject: [PATCH] py/emitnative: Add support for using setjmp with native
 emitter.

To enable this feature the N_NLR_SETJMP macro should be set to 1 before
including py/emitnative.c.
---
 py/emitnative.c | 14 ++++++++++++++
 py/nativeglue.c |  9 +++++++++
 py/runtime0.h   |  1 +
 3 files changed, 24 insertions(+)

diff --git a/py/emitnative.c b/py/emitnative.c
index 30f66f633..22bfa2c78 100644
--- a/py/emitnative.c
+++ b/py/emitnative.c
@@ -1154,6 +1154,10 @@ STATIC void emit_native_global_exc_entry(emit_t *emit) {
             // Wrap everything in an nlr context
             ASM_MOV_REG_LOCAL_ADDR(emit->as, REG_ARG_1, 0);
             emit_call(emit, MP_F_NLR_PUSH);
+            #if N_NLR_SETJMP
+            ASM_MOV_REG_LOCAL_ADDR(emit->as, REG_ARG_1, 2);
+            emit_call(emit, MP_F_SETJMP);
+            #endif
             ASM_JUMP_IF_REG_ZERO(emit->as, REG_RET, start_label, true);
         } else {
             // Clear the unwind state
@@ -1168,6 +1172,10 @@ STATIC void emit_native_global_exc_entry(emit_t *emit) {
             ASM_MOV_REG_LOCAL(emit->as, REG_LOCAL_2, LOCAL_IDX_EXC_HANDLER_UNWIND(emit));
             ASM_MOV_REG_LOCAL_ADDR(emit->as, REG_ARG_1, 0);
             emit_call(emit, MP_F_NLR_PUSH);
+            #if N_NLR_SETJMP
+            ASM_MOV_REG_LOCAL_ADDR(emit->as, REG_ARG_1, 2);
+            emit_call(emit, MP_F_SETJMP);
+            #endif
             ASM_MOV_LOCAL_REG(emit->as, LOCAL_IDX_EXC_HANDLER_UNWIND(emit), REG_LOCAL_2);
             ASM_JUMP_IF_REG_NONZERO(emit->as, REG_RET, global_except_label, true);
 
@@ -1178,6 +1186,12 @@ STATIC void emit_native_global_exc_entry(emit_t *emit) {
 
             // Global exception handler: check for valid exception handler
             emit_native_label_assign(emit, global_except_label);
+            #if N_NLR_SETJMP
+            // Reload REG_FUN_TABLE, since it may be clobbered by longjmp
+            emit_native_mov_reg_state(emit, REG_LOCAL_1, LOCAL_IDX_FUN_OBJ(emit));
+            ASM_LOAD_REG_REG_OFFSET(emit->as, REG_LOCAL_1, REG_LOCAL_1, offsetof(mp_obj_fun_bc_t, const_table) / sizeof(uintptr_t));
+            ASM_LOAD_REG_REG_OFFSET(emit->as, REG_FUN_TABLE, REG_LOCAL_1, emit->scope->num_pos_args + emit->scope->num_kwonly_args);
+            #endif
             ASM_MOV_REG_LOCAL(emit->as, REG_LOCAL_1, LOCAL_IDX_EXC_HANDLER_PC(emit));
             ASM_JUMP_IF_REG_NONZERO(emit->as, REG_LOCAL_1, nlr_label, false);
         }
diff --git a/py/nativeglue.c b/py/nativeglue.c
index 62b76eb6b..4405b2d11 100644
--- a/py/nativeglue.c
+++ b/py/nativeglue.c
@@ -244,7 +244,11 @@ const void *const mp_fun_table[MP_F_NUMBER_OF] = {
     mp_call_method_n_kw_var,
     mp_native_getiter,
     mp_native_iternext,
+    #if MICROPY_NLR_SETJMP
+    nlr_push_tail,
+    #else
     nlr_push,
+    #endif
     nlr_pop,
     mp_native_raise,
     mp_import_name,
@@ -262,6 +266,11 @@ const void *const mp_fun_table[MP_F_NUMBER_OF] = {
     mp_small_int_floor_divide,
     mp_small_int_modulo,
     mp_native_yield_from,
+    #if MICROPY_NLR_SETJMP
+    setjmp,
+    #else
+    NULL,
+    #endif
 };
 
 #endif // MICROPY_EMIT_NATIVE
diff --git a/py/runtime0.h b/py/runtime0.h
index 1df6d0d58..797ae00e6 100644
--- a/py/runtime0.h
+++ b/py/runtime0.h
@@ -201,6 +201,7 @@ typedef enum {
     MP_F_SMALL_INT_FLOOR_DIVIDE,
     MP_F_SMALL_INT_MODULO,
     MP_F_NATIVE_YIELD_FROM,
+    MP_F_SETJMP,
     MP_F_NUMBER_OF,
 } mp_fun_kind_t;