py/compile: Fix case of eager implicit conversion of local to nonlocal.

This ensures that implicit variables are only converted to implicit
closed-over variables (nonlocals) at the very end of the function scope.
If variables are closed-over when first used (read from, as was done prior
to this commit) then this can be incorrect because the variable may be
assigned to later on in the function which means they are just a plain
local, not closed over.

Fixes issue #4272.

4 files changed, 15 insertions, 8 deletions

diff --git a/py/compile.c b/py/compile.c
index e90b366e0..e708bde8e 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -1196,7 +1196,8 @@ STATIC void compile_declare_global(compiler_t *comp, mp_parse_node_t pn, qstr qs
 
 STATIC void compile_declare_nonlocal(compiler_t *comp, mp_parse_node_t pn, qstr qst, bool added, id_info_t *id_info) {
     if (added) {
-        scope_find_local_and_close_over(comp->scope_cur, id_info, qst);
+        id_info->kind = ID_INFO_KIND_GLOBAL_IMPLICIT;
+        scope_check_to_close_over(comp->scope_cur, id_info);
         if (id_info->kind == ID_INFO_KIND_GLOBAL_IMPLICIT) {
             compile_syntax_error(comp, pn, "no binding for nonlocal found");
         }
@@ -3391,6 +3392,14 @@ mp_raw_code_t *mp_compile_to_raw_code(mp_parse_tree_t *parse_tree, qstr source_f
         #endif
         } else {
             compile_scope(comp, s, MP_PASS_SCOPE);
+
+            // Check if any implicitly declared variables should be closed over
+            for (size_t i = 0; i < s->id_info_len; ++i) {
+                id_info_t *id = &s->id_info[i];
+                if (id->kind == ID_INFO_KIND_GLOBAL_IMPLICIT) {
+                    scope_check_to_close_over(s, id);
+                }
+            }
         }
 
         // update maximim number of labels needed
diff --git a/py/emitcommon.c b/py/emitcommon.c
index 89cc2c959..149e0b0f1 100644
--- a/py/emitcommon.c
+++ b/py/emitcommon.c
@@ -35,7 +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) {
-        scope_find_local_and_close_over(scope, id, qst);
+        id->kind = ID_INFO_KIND_GLOBAL_IMPLICIT;
     }
 }
 
diff --git a/py/scope.c b/py/scope.c
index 1a6ae7b8a..8adb85b80 100644
--- a/py/scope.c
+++ b/py/scope.c
@@ -130,21 +130,19 @@ STATIC void scope_close_over_in_parents(scope_t *scope, qstr qst) {
     }
 }
 
-void scope_find_local_and_close_over(scope_t *scope, id_info_t *id, qstr qst) {
+void scope_check_to_close_over(scope_t *scope, id_info_t *id) {
     if (scope->parent != NULL) {
         for (scope_t *s = scope->parent; s->parent != NULL; s = s->parent) {
-            id_info_t *id2 = scope_find(s, qst);
+            id_info_t *id2 = scope_find(s, id->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;
+                    scope_close_over_in_parents(scope, id->qst);
                 }
                 break;
             }
         }
     }
-    id->kind = ID_INFO_KIND_GLOBAL_IMPLICIT;
 }
 
 #endif // MICROPY_ENABLE_COMPILER
diff --git a/py/scope.h b/py/scope.h
index 5e9a0eb7b..d51bb90bb 100644
--- a/py/scope.h
+++ b/py/scope.h
@@ -93,6 +93,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);
-void scope_find_local_and_close_over(scope_t *scope, id_info_t *id, qstr qst);
+void scope_check_to_close_over(scope_t *scope, id_info_t *id);
 
 #endif // MICROPY_INCLUDED_PY_SCOPE_H