1 files changed, 43 insertions, 46 deletions

diff --git a/py/scope.c b/py/scope.c
index e408251733..8143e655c9 100644
--- a/py/scope.c
+++ b/py/scope.c
@@ -30,36 +30,27 @@
 
 #if MICROPY_ENABLE_COMPILER
 
+// these low numbered qstrs should fit in 8 bits
+STATIC const uint8_t scope_simple_name_table[] = {
+    [SCOPE_MODULE] = MP_QSTR__lt_module_gt_,
+    [SCOPE_LAMBDA] = MP_QSTR__lt_lambda_gt_,
+    [SCOPE_LIST_COMP] = MP_QSTR__lt_listcomp_gt_,
+    [SCOPE_DICT_COMP] = MP_QSTR__lt_dictcomp_gt_,
+    [SCOPE_SET_COMP] = MP_QSTR__lt_setcomp_gt_,
+    [SCOPE_GEN_EXPR] = MP_QSTR__lt_genexpr_gt_,
+};
+
 scope_t *scope_new(scope_kind_t kind, const byte *pn, qstr source_file, mp_uint_t emit_options) {
     scope_t *scope = m_new0(scope_t, 1);
     scope->kind = kind;
     scope->pn = pn;
     scope->source_file = source_file;
-    switch (kind) {
-        case SCOPE_MODULE:
-            scope->simple_name = MP_QSTR__lt_module_gt_;
-            break;
-        case SCOPE_FUNCTION:
-        case SCOPE_CLASS:
-            pt_extract_id(pn, &scope->simple_name); // function name
-            break;
-        case SCOPE_LAMBDA:
-            scope->simple_name = MP_QSTR__lt_lambda_gt_;
-            break;
-        case SCOPE_LIST_COMP:
-            scope->simple_name = MP_QSTR__lt_listcomp_gt_;
-            break;
-        case SCOPE_DICT_COMP:
-            scope->simple_name = MP_QSTR__lt_dictcomp_gt_;
-            break;
-        case SCOPE_SET_COMP:
-            scope->simple_name = MP_QSTR__lt_setcomp_gt_;
-            break;
-        case SCOPE_GEN_EXPR:
-            scope->simple_name = MP_QSTR__lt_genexpr_gt_;
-            break;
-        default:
-            assert(0);
+    if (kind == SCOPE_FUNCTION || kind == SCOPE_CLASS) {
+        qstr id;
+        pt_extract_id(pn, &id); // function name
+        scope->simple_name = id;
+    } else {
+        scope->simple_name = scope_simple_name_table[kind];
     }
     scope->emit_options = emit_options;
     scope->id_info_alloc = MICROPY_ALLOC_SCOPE_ID_INIT;
@@ -115,22 +106,10 @@ id_info_t *scope_find_global(scope_t *scope, qstr qst) {
     return scope_find(scope, qst);
 }
 
-id_info_t *scope_find_local_in_parent(scope_t *scope, qstr qst) {
-    if (scope->parent == NULL) {
-        return NULL;
-    }
-    for (scope_t *s = scope->parent; s->parent != NULL; s = s->parent) {
-        id_info_t *id = scope_find(s, qst);
-        if (id != NULL) {
-            return id;
-        }
-    }
-    return NULL;
-}
-
-void scope_close_over_in_parents(scope_t *scope, qstr qst) {
+STATIC void scope_close_over_in_parents(scope_t *scope, qstr qst) {
     assert(scope->parent != NULL); // we should have at least 1 parent
-    for (scope_t *s = scope->parent; s->parent != NULL; s = s->parent) {
+    for (scope_t *s = scope->parent;; s = s->parent) {
+        assert(s->parent != NULL); // we should not get to the outer scope
         bool added;
         id_info_t *id = scope_find_or_add_id(s, qst, &added);
         if (added) {
@@ -138,16 +117,34 @@ void scope_close_over_in_parents(scope_t *scope, qstr qst) {
             id->kind = ID_INFO_KIND_FREE;
         } else {
             // variable is declared in this scope, so finish
-            switch (id->kind) {
-                case ID_INFO_KIND_LOCAL: id->kind = ID_INFO_KIND_CELL; break; // variable local to this scope, close it over
-                case ID_INFO_KIND_FREE: break; // variable already closed over in a parent scope
-                case ID_INFO_KIND_CELL: break; // variable already closed over in this scope
-                default: assert(0); // TODO
+            if (id->kind == ID_INFO_KIND_LOCAL) {
+                // variable local to this scope, close it over
+                id->kind = ID_INFO_KIND_CELL;
+            } else {
+                // ID_INFO_KIND_FREE: variable already closed over in a parent scope
+                // ID_INFO_KIND_CELL: variable already closed over in this scope
+                assert(id->kind == ID_INFO_KIND_FREE || id->kind == ID_INFO_KIND_CELL);
             }
             return;
         }
     }
-    assert(0); // we should have found the variable in one of the parents
+}
+
+void scope_find_local_and_close_over(scope_t *scope, id_info_t *id, qstr qst) {
+    if (scope->parent != NULL) {
+        for (scope_t *s = scope->parent; s->parent != NULL; s = s->parent) {
+            id_info_t *id2 = scope_find(s, qst);
+            if (id2 != NULL) {
+                if (id2->kind == ID_INFO_KIND_LOCAL || id2->kind == ID_INFO_KIND_CELL || id2->kind == ID_INFO_KIND_FREE) {
+                    id->kind = ID_INFO_KIND_FREE;
+                    scope_close_over_in_parents(scope, qst);
+                    return;
+                }
+                break;
+            }
+        }
+    }
+    id->kind = ID_INFO_KIND_GLOBAL_IMPLICIT;
 }
 
 #endif // MICROPY_ENABLE_COMPILER