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

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

1 files changed, 20 insertions, 19 deletions

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;