py/compile: Fix scope of assignment expression target in comprehensions.

When := is used in a comprehension the target variable is bound to the
parent scope, so it's either a global or a nonlocal.  Prior to this commit
that was handled by simply using the parent scope's id_info for the
target variable.  That's completely wrong because it uses the slot number
for the parent's Python stack to store the variable, rather than the slot
number for the comprehension.  This will in most cases lead to incorrect
behaviour or memory faults.

This commit fixes the scoping of the target variable by explicitly
declaring it a global or nonlocal, depending on whether the parent is the
global scope or not.  Then the id_info of the comprehension can be used to
access the target variable.  This fixes a lot of cases of using := in a
comprehension.

Code size change for this commit:

       bare-arm:    +0 +0.000%
    minimal x86:    +0 +0.000%
       unix x64:  +152 +0.019% standard
          stm32:   +96 +0.024% PYBV10
         cc3200:   +96 +0.052%
        esp8266:  +196 +0.028% GENERIC
          esp32:  +156 +0.010% GENERIC[incl +8(data)]
         mimxrt:   +96 +0.027% TEENSY40
     renesas-ra:   +88 +0.014% RA6M2_EK
            nrf:   +88 +0.048% pca10040
            rp2:  +104 +0.020% PICO
           samd:   +88 +0.033% ADAFRUIT_ITSYBITSY_M4_EXPRESS

Fixes issue #10895.

Signed-off-by: Damien George <damien@micropython.org>

3 files changed, 35 insertions, 13 deletions

diff --git a/py/compile.c b/py/compile.c
index 76d4c1bf5..7a1660b1b 100644
--- a/py/compile.c
+++ b/py/compile.c
@@ -1274,6 +1274,14 @@ STATIC void compile_declare_nonlocal(compiler_t *comp, mp_parse_node_t pn, id_in
     }
 }
 
+STATIC void compile_declare_global_or_nonlocal(compiler_t *comp, mp_parse_node_t pn, id_info_t *id_info, bool is_global) {
+    if (is_global) {
+        compile_declare_global(comp, pn, id_info);
+    } else {
+        compile_declare_nonlocal(comp, pn, id_info);
+    }
+}
+
 STATIC void compile_global_nonlocal_stmt(compiler_t *comp, mp_parse_node_struct_t *pns) {
     if (comp->pass == MP_PASS_SCOPE) {
         bool is_global = MP_PARSE_NODE_STRUCT_KIND(pns) == PN_global_stmt;
@@ -1288,11 +1296,7 @@ STATIC void compile_global_nonlocal_stmt(compiler_t *comp, mp_parse_node_struct_
         for (size_t i = 0; i < n; i++) {
             qstr qst = MP_PARSE_NODE_LEAF_ARG(nodes[i]);
             id_info_t *id_info = scope_find_or_add_id(comp->scope_cur, qst, ID_INFO_KIND_UNDECIDED);
-            if (is_global) {
-                compile_declare_global(comp, (mp_parse_node_t)pns, id_info);
-            } else {
-                compile_declare_nonlocal(comp, (mp_parse_node_t)pns, id_info);
-            }
+            compile_declare_global_or_nonlocal(comp, (mp_parse_node_t)pns, id_info, is_global);
         }
     }
 }
@@ -2133,13 +2137,30 @@ STATIC void compile_namedexpr_helper(compiler_t *comp, mp_parse_node_t pn_name,
     }
     compile_node(comp, pn_expr);
     EMIT(dup_top);
-    scope_t *old_scope = comp->scope_cur;
-    if (SCOPE_IS_COMP_LIKE(comp->scope_cur->kind)) {
-        // Use parent's scope for assigned value so it can "escape"
-        comp->scope_cur = comp->scope_cur->parent;
+
+    qstr target = MP_PARSE_NODE_LEAF_ARG(pn_name);
+
+    // When a variable is assigned via := in a comprehension then that variable is bound to
+    // the parent scope.  Any global or nonlocal declarations in the parent scope are honoured.
+    // For details see: https://peps.python.org/pep-0572/#scope-of-the-target
+    if (comp->pass == MP_PASS_SCOPE && SCOPE_IS_COMP_LIKE(comp->scope_cur->kind)) {
+        id_info_t *id_info_parent = mp_emit_common_get_id_for_modification(comp->scope_cur->parent, target);
+        if (id_info_parent->kind == ID_INFO_KIND_GLOBAL_EXPLICIT) {
+            scope_find_or_add_id(comp->scope_cur, target, ID_INFO_KIND_GLOBAL_EXPLICIT);
+        } else {
+            id_info_t *id_info = scope_find_or_add_id(comp->scope_cur, target, ID_INFO_KIND_UNDECIDED);
+            bool is_global = comp->scope_cur->parent->parent == NULL; // comprehension is defined in outer scope
+            if (!is_global && id_info->kind == ID_INFO_KIND_GLOBAL_IMPLICIT) {
+                // Variable was already referenced but now needs to be closed over, so reset the kind
+                // such that scope_check_to_close_over() is called in compile_declare_nonlocal().
+                id_info->kind = ID_INFO_KIND_UNDECIDED;
+            }
+            compile_declare_global_or_nonlocal(comp, pn_name, id_info, is_global);
+        }
     }
-    compile_store_id(comp, MP_PARSE_NODE_LEAF_ARG(pn_name));
-    comp->scope_cur = old_scope;
+
+    // Do the store to the target variable.
+    compile_store_id(comp, target);
 }
 
 STATIC void compile_namedexpr(compiler_t *comp, mp_parse_node_struct_t *pns) {
diff --git a/py/emit.h b/py/emit.h
index 4e8a55e77..26f978ba5 100644
--- a/py/emit.h
+++ b/py/emit.h
@@ -191,7 +191,7 @@ static inline void mp_emit_common_get_id_for_load(scope_t *scope, qstr qst) {
     scope_find_or_add_id(scope, qst, ID_INFO_KIND_GLOBAL_IMPLICIT);
 }
 
-void mp_emit_common_get_id_for_modification(scope_t *scope, qstr qst);
+id_info_t *mp_emit_common_get_id_for_modification(scope_t *scope, qstr qst);
 void mp_emit_common_id_op(emit_t *emit, const mp_emit_method_table_id_ops_t *emit_method_table, scope_t *scope, qstr qst);
 
 extern const emit_method_table_t emit_bc_method_table;
diff --git a/py/emitcommon.c b/py/emitcommon.c
index 679ef1d97..a9eb6e202 100644
--- a/py/emitcommon.c
+++ b/py/emitcommon.c
@@ -86,7 +86,7 @@ size_t mp_emit_common_use_const_obj(mp_emit_common_t *emit, mp_obj_t const_obj)
     return emit->const_obj_list.len - 1;
 }
 
-void mp_emit_common_get_id_for_modification(scope_t *scope, qstr qst) {
+id_info_t *mp_emit_common_get_id_for_modification(scope_t *scope, qstr qst) {
     // name adding/lookup
     id_info_t *id = scope_find_or_add_id(scope, qst, ID_INFO_KIND_GLOBAL_IMPLICIT);
     if (id->kind == ID_INFO_KIND_GLOBAL_IMPLICIT) {
@@ -98,6 +98,7 @@ void mp_emit_common_get_id_for_modification(scope_t *scope, qstr qst) {
             id->kind = ID_INFO_KIND_GLOBAL_IMPLICIT_ASSIGNED;
         }
     }
+    return id;
 }
 
 void mp_emit_common_id_op(emit_t *emit, const mp_emit_method_table_id_ops_t *emit_method_table, scope_t *scope, qstr qst) {