py: Clear state to MP_OBJ_NULL before executing byte code.

2 files changed, 12 insertions, 10 deletions

diff --git a/py/objgenerator.c b/py/objgenerator.c
index 3a7f9f9d40..6468aa20db 100644
--- a/py/objgenerator.c
+++ b/py/objgenerator.c
@@ -252,14 +252,15 @@ mp_obj_t mp_obj_new_gen_instance(const byte *bytecode, uint n_args, const mp_obj
         o->state[n_state - 1 - n_args - i] = args2[i];
     }
 
+    // set rest of state to MP_OBJ_NULL
+    for (uint i = 0; i < n_state - n_args - n_args2; i++) {
+        o->state[i] = MP_OBJ_NULL;
+    }
+
     // bytecode prelude: initialise closed over variables
     for (uint n_local = *bytecode++; n_local > 0; n_local--) {
         uint local_num = *bytecode++;
-        if (local_num < n_args + n_args2) {
-            o->state[n_state - 1 - local_num] = mp_obj_new_cell(o->state[n_state - 1 - local_num]);
-        } else {
-            o->state[n_state - 1 - local_num] = mp_obj_new_cell(MP_OBJ_NULL);
-        }
+        o->state[n_state - 1 - local_num] = mp_obj_new_cell(o->state[n_state - 1 - local_num]);
     }
 
     // set ip to start of actual byte code
diff --git a/py/vm.c b/py/vm.c
index f8b60acd22..b9147d1d32 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -95,14 +95,15 @@ mp_vm_return_kind_t mp_execute_byte_code(const byte *code, const mp_obj_t *args,
         state[n_state - 1 - n_args - i] = args2[i];
     }
 
+    // set rest of state to MP_OBJ_NULL
+    for (uint i = 0; i < n_state - n_args - n_args2; i++) {
+        state[i] = MP_OBJ_NULL;
+    }
+
     // bytecode prelude: initialise closed over variables
     for (uint n_local = *ip++; n_local > 0; n_local--) {
         uint local_num = *ip++;
-        if (local_num < n_args + n_args2) {
-            state[n_state - 1 - local_num] = mp_obj_new_cell(state[n_state - 1 - local_num]);
-        } else {
-            state[n_state - 1 - local_num] = mp_obj_new_cell(MP_OBJ_NULL);
-        }
+        state[n_state - 1 - local_num] = mp_obj_new_cell(state[n_state - 1 - local_num]);
     }
 
     // execute the byte code