py: make closures work.

1 files changed, 15 insertions, 5 deletions

diff --git a/py/vm.c b/py/vm.c
index a23b0e89fb..3c3f398c04 100644
--- a/py/vm.c
+++ b/py/vm.c
@@ -39,10 +39,25 @@ mp_obj_t mp_execute_byte_code(const byte *code, const mp_obj_t *args, uint n_arg
         state[i] = args[n_args - 1 - i];
     }
     const byte *ip = code;
+
+    // execute prelude to make any cells (closed over variables)
+    {
+        for (uint n_local = *ip++; n_local > 0; n_local--) {
+            uint local_num = *ip++;
+            if (local_num < n_args) {
+                state[local_num] = mp_obj_new_cell(state[local_num]);
+            } else {
+                state[local_num] = mp_obj_new_cell(MP_OBJ_NULL);
+            }
+        }
+    }
+
+    // execute the byte code
     if (mp_execute_byte_code_2(&ip, &state[0], &sp)) {
         // it shouldn't yield
         assert(0);
     }
+
     // TODO check fails if, eg, return from within for loop
     //assert(sp == &state[17]);
     return *sp;
@@ -127,11 +142,6 @@ bool mp_execute_byte_code_2(const byte **ip_in_out, mp_obj_t *fastn, mp_obj_t **
                         PUSH(rt_get_cell(fastn[unum]));
                         break;
 
-                    case MP_BC_LOAD_CLOSURE:
-                        DECODE_UINT;
-                        PUSH(fastn[unum]);
-                        break;
-
                     case MP_BC_LOAD_NAME:
                         DECODE_QSTR;
                         PUSH(rt_load_name(qstr));