gh-104602: ensure all cellvars are known up front (#104603)

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>

4 files changed, 50 insertions, 22 deletions

diff --git a/Include/internal/pycore_symtable.h b/Include/internal/pycore_symtable.h
index 3fa825d..c8e0578 100644
--- a/Include/internal/pycore_symtable.h
+++ b/Include/internal/pycore_symtable.h
@@ -120,14 +120,15 @@ extern PyObject* _Py_Mangle(PyObject *p, PyObject *name);
 #define DEF_ANNOT 2<<7         /* this name is annotated */
 #define DEF_COMP_ITER 2<<8     /* this name is a comprehension iteration variable */
 #define DEF_TYPE_PARAM 2<<9    /* this name is a type parameter */
+#define DEF_COMP_CELL 2<<10    /* this name is a cell in an inlined comprehension */
 
 #define DEF_BOUND (DEF_LOCAL | DEF_PARAM | DEF_IMPORT)
 
 /* GLOBAL_EXPLICIT and GLOBAL_IMPLICIT are used internally by the symbol
    table.  GLOBAL is returned from PyST_GetScope() for either of them.
-   It is stored in ste_symbols at bits 12-15.
+   It is stored in ste_symbols at bits 13-16.
 */
-#define SCOPE_OFFSET 11
+#define SCOPE_OFFSET 12
 #define SCOPE_MASK (DEF_GLOBAL | DEF_LOCAL | DEF_PARAM | DEF_NONLOCAL)
 
 #define LOCAL 1
diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
index ae63f4a..fdd2d66 100644
--- a/Lib/test/test_listcomps.py
+++ b/Lib/test/test_listcomps.py
@@ -381,6 +381,32 @@ class ListComprehensionTest(unittest.TestCase):
         with self.assertRaises(UnboundLocalError):
             f()
 
+    def test_global_outside_cellvar_inside_plus_freevar(self):
+        code = """
+            a = 1
+            def f():
+                func, = [(lambda: b) for b in [a]]
+                return b, func()
+            x = f()
+        """
+        self._check_in_scopes(
+            code, {"x": (2, 1)}, ns={"b": 2}, scopes=["function", "module"])
+        # inside a class, the `a = 1` assignment is not visible
+        self._check_in_scopes(code, raises=NameError, scopes=["class"])
+
+    def test_cell_in_nested_comprehension(self):
+        code = """
+            a = 1
+            def f():
+                (func, inner_b), = [[lambda: b for b in c] + [b] for c in [[a]]]
+                return b, inner_b, func()
+            x = f()
+        """
+        self._check_in_scopes(
+            code, {"x": (2, 2, 1)}, ns={"b": 2}, scopes=["function", "module"])
+        # inside a class, the `a = 1` assignment is not visible
+        self._check_in_scopes(code, raises=NameError, scopes=["class"])
+
     def test_name_error_in_class_scope(self):
         code = """
             y = 1
diff --git a/Python/compile.c b/Python/compile.c
index 96cd6da..2db03f7 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -1252,7 +1252,7 @@ compiler_enter_scope(struct compiler *c, identifier name,
     }
     u->u_metadata.u_name = Py_NewRef(name);
     u->u_metadata.u_varnames = list2dict(u->u_ste->ste_varnames);
-    u->u_metadata.u_cellvars = dictbytype(u->u_ste->ste_symbols, CELL, 0, 0);
+    u->u_metadata.u_cellvars = dictbytype(u->u_ste->ste_symbols, CELL, DEF_COMP_CELL, 0);
     if (!u->u_metadata.u_varnames || !u->u_metadata.u_cellvars) {
         compiler_unit_free(u);
         return ERROR;
diff --git a/Python/symtable.c b/Python/symtable.c
index a319c23..bd523f0 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -632,7 +632,7 @@ is_free_in_any_child(PySTEntryObject *entry, PyObject *key)
 static int
 inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
                      PyObject *scopes, PyObject *comp_free,
-                     PyObject *promote_to_cell)
+                     PyObject *inlined_cells)
 {
     PyObject *k, *v;
     Py_ssize_t pos = 0;
@@ -645,6 +645,11 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
         }
         int scope = (comp_flags >> SCOPE_OFFSET) & SCOPE_MASK;
         int only_flags = comp_flags & ((1 << SCOPE_OFFSET) - 1);
+        if (scope == CELL || only_flags & DEF_COMP_CELL) {
+            if (PySet_Add(inlined_cells, k) < 0) {
+                return 0;
+            }
+        }
         PyObject *existing = PyDict_GetItemWithError(ste->ste_symbols, k);
         if (existing == NULL && PyErr_Occurred()) {
             return 0;
@@ -665,14 +670,6 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
         }
         else {
             if (PyLong_AsLong(existing) & DEF_BOUND) {
-                // cell vars in comprehension that are locals in outer scope
-                // must be promoted to cell so u_cellvars isn't wrong
-                if (scope == CELL && _PyST_IsFunctionLike(ste)) {
-                    if (PySet_Add(promote_to_cell, k) < 0) {
-                        return 0;
-                    }
-                }
-
                 // free vars in comprehension that are locals in outer scope can
                 // now simply be locals, unless they are free in comp children,
                 // or if the outer scope is a class block
@@ -698,7 +695,7 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
 */
 
 static int
-analyze_cells(PyObject *scopes, PyObject *free, PyObject *promote_to_cell)
+analyze_cells(PyObject *scopes, PyObject *free, PyObject *inlined_cells)
 {
     PyObject *name, *v, *v_cell;
     int success = 0;
@@ -713,7 +710,7 @@ analyze_cells(PyObject *scopes, PyObject *free, PyObject *promote_to_cell)
         scope = PyLong_AS_LONG(v);
         if (scope != LOCAL)
             continue;
-        if (!PySet_Contains(free, name) && !PySet_Contains(promote_to_cell, name))
+        if (!PySet_Contains(free, name) && !PySet_Contains(inlined_cells, name))
             continue;
         /* Replace LOCAL with CELL for this name, and remove
            from free. It is safe to replace the value of name
@@ -753,7 +750,8 @@ drop_class_free(PySTEntryObject *ste, PyObject *free)
 */
 static int
 update_symbols(PyObject *symbols, PyObject *scopes,
-               PyObject *bound, PyObject *free, int classflag)
+               PyObject *bound, PyObject *free,
+               PyObject *inlined_cells, int classflag)
 {
     PyObject *name = NULL, *itr = NULL;
     PyObject *v = NULL, *v_scope = NULL, *v_new = NULL, *v_free = NULL;
@@ -764,6 +762,9 @@ update_symbols(PyObject *symbols, PyObject *scopes,
         long scope, flags;
         assert(PyLong_Check(v));
         flags = PyLong_AS_LONG(v);
+        if (PySet_Contains(inlined_cells, name)) {
+            flags |= DEF_COMP_CELL;
+        }
         v_scope = PyDict_GetItemWithError(scopes, name);
         assert(v_scope && PyLong_Check(v_scope));
         scope = PyLong_AS_LONG(v_scope);
@@ -870,7 +871,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
               PySTEntryObject *class_entry)
 {
     PyObject *name, *v, *local = NULL, *scopes = NULL, *newbound = NULL;
-    PyObject *newglobal = NULL, *newfree = NULL, *promote_to_cell = NULL;
+    PyObject *newglobal = NULL, *newfree = NULL, *inlined_cells = NULL;
     PyObject *temp;
     int success = 0;
     Py_ssize_t i, pos = 0;
@@ -902,8 +903,8 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
     newbound = PySet_New(NULL);
     if (!newbound)
         goto error;
-    promote_to_cell = PySet_New(NULL);
-    if (!promote_to_cell)
+    inlined_cells = PySet_New(NULL);
+    if (!inlined_cells)
         goto error;
 
     /* Class namespace has no effect on names visible in
@@ -997,7 +998,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
             goto error;
         }
         if (inline_comp) {
-            if (!inline_comprehension(ste, entry, scopes, child_free, promote_to_cell)) {
+            if (!inline_comprehension(ste, entry, scopes, child_free, inlined_cells)) {
                 Py_DECREF(child_free);
                 goto error;
             }
@@ -1028,12 +1029,12 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
     }
 
     /* Check if any local variables must be converted to cell variables */
-    if (_PyST_IsFunctionLike(ste) && !analyze_cells(scopes, newfree, promote_to_cell))
+    if (_PyST_IsFunctionLike(ste) && !analyze_cells(scopes, newfree, inlined_cells))
         goto error;
     else if (ste->ste_type == ClassBlock && !drop_class_free(ste, newfree))
         goto error;
     /* Records the results of the analysis in the symbol table entry */
-    if (!update_symbols(ste->ste_symbols, scopes, bound, newfree,
+    if (!update_symbols(ste->ste_symbols, scopes, bound, newfree, inlined_cells,
                         ste->ste_type == ClassBlock))
         goto error;
 
@@ -1048,7 +1049,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
     Py_XDECREF(newbound);
     Py_XDECREF(newglobal);
     Py_XDECREF(newfree);
-    Py_XDECREF(promote_to_cell);
+    Py_XDECREF(inlined_cells);
     if (!success)
         assert(PyErr_Occurred());
     return success;