py/scope: Factor common code to find locals and close over them.

Saves 50-100 bytes of code.

4 files changed, 35 insertions, 39 deletions

diff --git a/py/compile.c b/py/compile.c
index 2253ccd3f6..1d371cb330 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -1172,17 +1172,14 @@ STATIC void compile_global_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
 STATIC void compile_declare_nonlocal(compiler_t *comp, mp_parse_node_t pn, qstr qst) {
     bool added;
     id_info_t *id_info = scope_find_or_add_id(comp->scope_cur, qst, &added);
-    if (!added && id_info->kind != ID_INFO_KIND_FREE) {
+    if (added) {
+        scope_find_local_and_close_over(comp->scope_cur, id_info, qst);
+        if (id_info->kind == ID_INFO_KIND_GLOBAL_IMPLICIT) {
+            compile_syntax_error(comp, pn, "no binding for nonlocal found");
+        }
+    } else if (id_info->kind != ID_INFO_KIND_FREE) {
         compile_syntax_error(comp, pn, "identifier redefined as nonlocal");
-        return;
-    }
-    id_info_t *id_info2 = scope_find_local_in_parent(comp->scope_cur, qst);
-    if (id_info2 == NULL || !(id_info2->kind == ID_INFO_KIND_LOCAL || id_info2->kind == ID_INFO_KIND_CELL || id_info2->kind == ID_INFO_KIND_FREE)) {
-        compile_syntax_error(comp, pn, "no binding for nonlocal found");
-        return;
     }
-    id_info->kind = ID_INFO_KIND_FREE;
-    scope_close_over_in_parents(comp->scope_cur, qst);
 }
 
 STATIC void compile_nonlocal_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
diff --git a/py/emitcommon.c b/py/emitcommon.c
index 2925f4cafe..e914431d32 100644
--- a/py/emitcommon.c
+++ b/py/emitcommon.c
@@ -35,13 +35,7 @@ void mp_emit_common_get_id_for_load(scope_t *scope, qstr qst) {
     bool added;
     id_info_t *id = scope_find_or_add_id(scope, qst, &added);
     if (added) {
-        id_info_t *id2 = scope_find_local_in_parent(scope, qst);
-        if (id2 != NULL && (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);
-        } else {
-            id->kind = ID_INFO_KIND_GLOBAL_IMPLICIT;
-        }
+        scope_find_local_and_close_over(scope, id, qst);
     }
 }
 
diff --git a/py/scope.c b/py/scope.c
index d080fbcbfb..8fe6f960ad 100644
--- a/py/scope.c
+++ b/py/scope.c
@@ -106,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) {
@@ -129,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
diff --git a/py/scope.h b/py/scope.h
index 52cbbf118f..826064d7ef 100644
--- a/py/scope.h
+++ b/py/scope.h
@@ -92,7 +92,6 @@ void scope_free(scope_t *scope);
 id_info_t *scope_find_or_add_id(scope_t *scope, qstr qstr, bool *added);
 id_info_t *scope_find(scope_t *scope, qstr qstr);
 id_info_t *scope_find_global(scope_t *scope, qstr qstr);
-id_info_t *scope_find_local_in_parent(scope_t *scope, qstr qstr);
-void scope_close_over_in_parents(scope_t *scope, qstr qstr);
+void scope_find_local_and_close_over(scope_t *scope, id_info_t *id, qstr qst);
 
 #endif // __MICROPY_INCLUDED_PY_SCOPE_H__