gh-104404: fix crasher with nested comprehensions plus lambdas (#104442)

carljm · web-flow · commit 563c7dcba0ea · 2023-05-12T17:42:04.000-07:00
diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
@@ -338,6 +338,14 @@ def test_nested_3(self):
         outputs = {"y": [1, 3, 5]}
         self._check_in_scopes(code, outputs)
 
+    def test_nested_4(self):
+        code = """
+            items = [([lambda: x for x in range(2)], lambda: x) for x in range(3)]
+            out = [([fn() for fn in fns], fn()) for fns, fn in items]
+        """
+        outputs = {"out": [([1, 1], 2), ([1, 1], 2), ([1, 1], 2)]}
+        self._check_in_scopes(code, outputs)
+
     def test_nameerror(self):
         code = """
             [x for x in [1]]
diff --git a/Python/symtable.c b/Python/symtable.c
@@ -575,7 +575,8 @@ is_free_in_any_child(PySTEntryObject *entry, PyObject *key)
 
 static int
 inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
-                     PyObject *scopes, PyObject *comp_free)
+                     PyObject *scopes, PyObject *comp_free,
+                     PyObject *promote_to_cell)
 {
     PyObject *k, *v;
     Py_ssize_t pos = 0;
@@ -611,7 +612,9 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
                 // cell vars in comprehension that are locals in outer scope
                 // must be promoted to cell so u_cellvars isn't wrong
                 if (scope == CELL && ste->ste_type == FunctionBlock) {
-                    SET_SCOPE(scopes, k, scope);
+                    if (PySet_Add(promote_to_cell, k) < 0) {
+                        return 0;
+                    }
                 }
 
                 // free vars in comprehension that are locals in outer scope can
@@ -638,7 +641,7 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
 */
 
 static int
-analyze_cells(PyObject *scopes, PyObject *free)
+analyze_cells(PyObject *scopes, PyObject *free, PyObject *promote_to_cell)
 {
     PyObject *name, *v, *v_cell;
     int success = 0;
@@ -653,7 +656,7 @@ analyze_cells(PyObject *scopes, PyObject *free)
         scope = PyLong_AS_LONG(v);
         if (scope != LOCAL)
             continue;
-        if (!PySet_Contains(free, name))
+        if (!PySet_Contains(free, name) && !PySet_Contains(promote_to_cell, name))
             continue;
         /* Replace LOCAL with CELL for this name, and remove
            from free. It is safe to replace the value of name
@@ -803,7 +806,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
               PyObject *global)
 {
     PyObject *name, *v, *local = NULL, *scopes = NULL, *newbound = NULL;
-    PyObject *newglobal = NULL, *newfree = NULL;
+    PyObject *newglobal = NULL, *newfree = NULL, *promote_to_cell = NULL;
     PyObject *temp;
     int success = 0;
     Py_ssize_t i, pos = 0;
@@ -835,6 +838,9 @@ 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)
+        goto error;
 
     /* Class namespace has no effect on names visible in
        nested functions, so populate the global and bound
@@ -915,7 +921,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
             goto error;
         }
         if (inline_comp) {
-            if (!inline_comprehension(ste, entry, scopes, child_free)) {
+            if (!inline_comprehension(ste, entry, scopes, child_free, promote_to_cell)) {
                 Py_DECREF(child_free);
                 goto error;
             }
@@ -946,7 +952,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
     }
 
     /* Check if any local variables must be converted to cell variables */
-    if (ste->ste_type == FunctionBlock && !analyze_cells(scopes, newfree))
+    if (ste->ste_type == FunctionBlock && !analyze_cells(scopes, newfree, promote_to_cell))
         goto error;
     else if (ste->ste_type == ClassBlock && !drop_class_free(ste, newfree))
         goto error;
@@ -966,6 +972,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
     Py_XDECREF(newbound);
     Py_XDECREF(newglobal);
     Py_XDECREF(newfree);
+    Py_XDECREF(promote_to_cell);
     if (!success)
         assert(PyErr_Occurred());
     return success;
