py: Add n_state to mp_code_state_t struct.

This value is used often enough that it is better to cache it instead of
decode it each time.

5 files changed, 12 insertions, 5 deletions

diff --git a/py/bc.h b/py/bc.h
index 1d28dcc4f9..69bce89028 100644
--- a/py/bc.h
+++ b/py/bc.h
@@ -99,6 +99,7 @@ typedef struct _mp_code_state_t {
     mp_obj_fun_bc_t *fun_bc;
     const byte *ip;
     mp_obj_t *sp;
+    uint16_t n_state;
     uint16_t exc_sp_idx;
     mp_obj_dict_t *old_globals;
     #if MICROPY_STACKLESS
diff --git a/py/emitnative.c b/py/emitnative.c
index 760c7fb0c1..f5a18c987a 100644
--- a/py/emitnative.c
+++ b/py/emitnative.c
@@ -521,6 +521,9 @@ STATIC void emit_native_start_pass(emit_t *emit, pass_kind_t pass, scope_t *scop
             // TODO this encoding may change size in the final pass, need to make it fixed
             emit_native_mov_state_imm_via(emit, emit->code_state_start + OFFSETOF_CODE_STATE_IP, emit->prelude_offset, REG_ARG_1);
 
+            // Set code_state.n_state (only works on little endian targets due to n_state being uint16_t)
+            emit_native_mov_state_imm_via(emit, emit->code_state_start + offsetof(mp_code_state_t, n_state) / sizeof(uintptr_t), emit->n_state, REG_ARG_1);
+
             // Put address of code_state into first arg
             ASM_MOV_REG_LOCAL_ADDR(emit->as, REG_ARG_1, emit->code_state_start);
 
diff --git a/py/objfun.c b/py/objfun.c
index e0c6fb9271..114367b4e0 100644
--- a/py/objfun.c
+++ b/py/objfun.c
@@ -206,9 +206,10 @@ STATIC void dump_args(const mp_obj_t *a, size_t sz) {
                            + n_exc_stack * sizeof(mp_exc_stack_t);                \
     }
 
-#define INIT_CODESTATE(code_state, _fun_bc, n_args, n_kw, args) \
+#define INIT_CODESTATE(code_state, _fun_bc, _n_state, n_args, n_kw, args) \
     code_state->fun_bc = _fun_bc; \
     code_state->ip = 0; \
+    code_state->n_state = _n_state; \
     mp_setup_code_state(code_state, n_args, n_kw, args); \
     code_state->old_globals = mp_globals_get();
 
@@ -235,7 +236,7 @@ mp_code_state_t *mp_obj_fun_bc_prepare_codestate(mp_obj_t self_in, size_t n_args
     }
     #endif
 
-    INIT_CODESTATE(code_state, self, n_args, n_kw, args);
+    INIT_CODESTATE(code_state, self, n_state, n_args, n_kw, args);
 
     // execute the byte code with the correct globals context
     mp_globals_set(self->globals);
@@ -280,7 +281,7 @@ STATIC mp_obj_t fun_bc_call(mp_obj_t self_in, size_t n_args, size_t n_kw, const
     }
     #endif
 
-    INIT_CODESTATE(code_state, self, n_args, n_kw, args);
+    INIT_CODESTATE(code_state, self, n_state, n_args, n_kw, args);
 
     // execute the byte code with the correct globals context
     mp_globals_set(self->globals);
diff --git a/py/objgenerator.c b/py/objgenerator.c
index d32a74f3fb..39dcbefb92 100644
--- a/py/objgenerator.c
+++ b/py/objgenerator.c
@@ -62,6 +62,7 @@ STATIC mp_obj_t gen_wrap_call(mp_obj_t self_in, size_t n_args, size_t n_kw, cons
     o->globals = self_fun->globals;
     o->code_state.fun_bc = self_fun;
     o->code_state.ip = 0;
+    o->code_state.n_state = n_state;
     mp_setup_code_state(&o->code_state, n_args, n_kw, args);
     return MP_OBJ_FROM_PTR(o);
 }
@@ -99,6 +100,7 @@ STATIC mp_obj_t native_gen_wrap_call(mp_obj_t self_in, size_t n_args, size_t n_k
     o->globals = self_fun->globals;
     o->code_state.fun_bc = self_fun;
     o->code_state.ip = (const byte*)prelude_offset;
+    o->code_state.n_state = n_state;
     mp_setup_code_state(&o->code_state, n_args, n_kw, args);
 
     // Indicate we are a native function, which doesn't use this variable
diff --git a/py/vm.c b/py/vm.c
index c0cd9ffbf8..d59771b6e8 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -228,7 +228,7 @@ FRAME_SETUP();
     mp_obj_t * /*const*/ fastn;
     mp_exc_stack_t * /*const*/ exc_stack;
     {
-        size_t n_state = mp_decode_uint_value(code_state->fun_bc->bytecode);
+        size_t n_state = code_state->n_state;
         fastn = &code_state->state[n_state - 1];
         exc_stack = (mp_exc_stack_t*)(code_state->state + n_state);
     }
@@ -1499,7 +1499,7 @@ unwind_loop:
                 mp_nonlocal_free(code_state, sizeof(mp_code_state_t));
                 #endif
                 code_state = new_code_state;
-                size_t n_state = mp_decode_uint_value(code_state->fun_bc->bytecode);
+                size_t n_state = code_state->n_state;
                 fastn = &code_state->state[n_state - 1];
                 exc_stack = (mp_exc_stack_t*)(code_state->state + n_state);
                 // variables that are visible to the exception handler (declared volatile)