fix u_fasthidden in nested case

carljm · carljm · commit 46c7a4f4eb61 · 2023-05-05T11:35:29.000-07:00
diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
@@ -249,6 +249,23 @@ def test_nested(self):
         outputs = {"y": [[0, 1], [0, 1, 4]]}
         self._check_in_scopes(code, outputs)
 
+    def test_nested_2(self):
+        code = """
+            l = [1, 2, 3]
+            x = 3
+            y = [x for [x ** x for x in range(x)][x - 1] in l]
+        """
+        outputs = {"y": [3, 3, 3]}
+        self._check_in_scopes(code, outputs)
+
+    def test_nested_3(self):
+        code = """
+            l = [(1, 2), (3, 4), (5, 6)]
+            y = [x for (x, [x ** x for x in range(x)][x - 1]) in l]
+        """
+        outputs = {"y": [1, 3, 5]}
+        self._check_in_scopes(code, outputs)
+
     def test_nameerror(self):
         code = """
             [x for x in [1]]
diff --git a/Python/compile.c b/Python/compile.c
@@ -4990,6 +4990,7 @@ compiler_async_comprehension_generator(struct compiler *c, location loc,
 typedef struct {
     PyObject *pushed_locals;
     PyObject *temp_symbols;
+    PyObject *fast_hidden;
 } inlined_comprehension_state;
 
 static int
@@ -5011,8 +5012,20 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
         if (symbol & DEF_LOCAL && ~symbol & DEF_NONLOCAL) {
             if (c->u->u_ste->ste_type != FunctionBlock) {
                 // non-function scope: override this name to use fast locals
-                if (PyDict_SetItem(c->u->u_metadata.u_fasthidden, k, Py_True) < 0) {
-                    return ERROR;
+                PyObject *orig = PyDict_GetItem(c->u->u_metadata.u_fasthidden, k);
+                if (orig != Py_True) {
+                    if (PyDict_SetItem(c->u->u_metadata.u_fasthidden, k, Py_True) < 0) {
+                        return ERROR;
+                    }
+                    if (state->fast_hidden == NULL) {
+                        state->fast_hidden = PySet_New(NULL);
+                        if (state->fast_hidden == NULL) {
+                            return ERROR;
+                        }
+                    }
+                    if (PySet_Add(state->fast_hidden, k) < 0) {
+                        return ERROR;
+                    }
                 }
             }
             long scope = (symbol >> SCOPE_OFFSET) & SCOPE_MASK;
@@ -5092,6 +5105,7 @@ pop_inlined_comprehension_state(struct compiler *c, location loc,
                 return ERROR;
             }
         }
+        Py_CLEAR(state.temp_symbols);
     }
     if (state.pushed_locals) {
         // pop names we pushed to stack earlier
@@ -5108,16 +5122,21 @@ pop_inlined_comprehension_state(struct compiler *c, location loc,
             }
             ADDOP_NAME(c, loc, STORE_FAST_MAYBE_NULL, k, varnames);
         }
+        Py_CLEAR(state.pushed_locals);
     }
-    pos = 0;
-    while (PyDict_Next(c->u->u_metadata.u_fasthidden, &pos, &k, &v)) {
-        if (v == Py_True) {
+    if (state.fast_hidden) {
+        while (PySet_Size(state.fast_hidden) > 0) {
+            PyObject *k = PySet_Pop(state.fast_hidden);
+            if (k == NULL) {
+                return ERROR;
+            }
             // we set to False instead of clearing, so we can track which names
             // were temporarily fast-locals and should use CO_FAST_HIDDEN
             if (PyDict_SetItem(c->u->u_metadata.u_fasthidden, k, Py_False)) {
                 return ERROR;
             }
         }
+        Py_CLEAR(state.fast_hidden);
     }
     return SUCCESS;
 }
@@ -5214,8 +5233,6 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
         if (pop_inlined_comprehension_state(c, loc, inline_state)) {
             goto error;
         }
-        Py_XDECREF(inline_state.pushed_locals);
-        Py_XDECREF(inline_state.temp_symbols);
         return SUCCESS;
     }
 
@@ -5265,6 +5282,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
     Py_XDECREF(entry);
     Py_XDECREF(inline_state.pushed_locals);
     Py_XDECREF(inline_state.temp_symbols);
+    Py_XDECREF(inline_state.fast_hidden);
     return ERROR;
 }
 

Original file line number	Diff line number	Diff line change
`@@ -4990,6 +4990,7 @@ compiler_async_comprehension_generator(struct compiler *c, location loc,`
`4990`	`4990`	`typedef struct {`
`4991`	`4991`	`PyObject *pushed_locals;`
`4992`	`4992`	`PyObject *temp_symbols;`
	`4993`	`+ PyObject *fast_hidden;`
`4993`	`4994`	`} inlined_comprehension_state;`
`4994`	`4995`
`4995`	`4996`	`static int`
`@@ -5011,8 +5012,20 @@ push_inlined_comprehension_state(struct compiler *c, location loc,`
`5011`	`5012`	`if (symbol & DEF_LOCAL && ~symbol & DEF_NONLOCAL) {`
`5012`	`5013`	`if (c->u->u_ste->ste_type != FunctionBlock) {`
`5013`	`5014`	`// non-function scope: override this name to use fast locals`
`5014`		`- if (PyDict_SetItem(c->u->u_metadata.u_fasthidden, k, Py_True) < 0) {`
`5015`		`- return ERROR;`
	`5015`	`+ PyObject *orig = PyDict_GetItem(c->u->u_metadata.u_fasthidden, k);`
	`5016`	`+ if (orig != Py_True) {`
	`5017`	`+ if (PyDict_SetItem(c->u->u_metadata.u_fasthidden, k, Py_True) < 0) {`
	`5018`	`+ return ERROR;`
	`5019`	`+ }`
	`5020`	`+ if (state->fast_hidden == NULL) {`
	`5021`	`+ state->fast_hidden = PySet_New(NULL);`
	`5022`	`+ if (state->fast_hidden == NULL) {`
	`5023`	`+ return ERROR;`
	`5024`	`+ }`
	`5025`	`+ }`
	`5026`	`+ if (PySet_Add(state->fast_hidden, k) < 0) {`
	`5027`	`+ return ERROR;`
	`5028`	`+ }`
`5016`	`5029`	`}`
`5017`	`5030`	`}`
`5018`	`5031`	`long scope = (symbol >> SCOPE_OFFSET) & SCOPE_MASK;`
`@@ -5092,6 +5105,7 @@ pop_inlined_comprehension_state(struct compiler *c, location loc,`
`5092`	`5105`	`return ERROR;`
`5093`	`5106`	`}`
`5094`	`5107`	`}`
	`5108`	`+ Py_CLEAR(state.temp_symbols);`
`5095`	`5109`	`}`
`5096`	`5110`	`if (state.pushed_locals) {`
`5097`	`5111`	`// pop names we pushed to stack earlier`
`@@ -5108,16 +5122,21 @@ pop_inlined_comprehension_state(struct compiler *c, location loc,`
`5108`	`5122`	`}`
`5109`	`5123`	`ADDOP_NAME(c, loc, STORE_FAST_MAYBE_NULL, k, varnames);`
`5110`	`5124`	`}`
	`5125`	`+ Py_CLEAR(state.pushed_locals);`
`5111`	`5126`	`}`
`5112`		`- pos = 0;`
`5113`		`- while (PyDict_Next(c->u->u_metadata.u_fasthidden, &pos, &k, &v)) {`
`5114`		`- if (v == Py_True) {`
	`5127`	`+ if (state.fast_hidden) {`
	`5128`	`+ while (PySet_Size(state.fast_hidden) > 0) {`
	`5129`	`+ PyObject *k = PySet_Pop(state.fast_hidden);`
	`5130`	`+ if (k == NULL) {`
	`5131`	`+ return ERROR;`
	`5132`	`+ }`
`5115`	`5133`	`// we set to False instead of clearing, so we can track which names`
`5116`	`5134`	`// were temporarily fast-locals and should use CO_FAST_HIDDEN`
`5117`	`5135`	`if (PyDict_SetItem(c->u->u_metadata.u_fasthidden, k, Py_False)) {`
`5118`	`5136`	`return ERROR;`
`5119`	`5137`	`}`
`5120`	`5138`	`}`
	`5139`	`+ Py_CLEAR(state.fast_hidden);`
`5121`	`5140`	`}`
`5122`	`5141`	`return SUCCESS;`
`5123`	`5142`	`}`
`@@ -5214,8 +5233,6 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,`
`5214`	`5233`	`if (pop_inlined_comprehension_state(c, loc, inline_state)) {`
`5215`	`5234`	`goto error;`
`5216`	`5235`	`}`
`5217`		`- Py_XDECREF(inline_state.pushed_locals);`
`5218`		`- Py_XDECREF(inline_state.temp_symbols);`
`5219`	`5236`	`return SUCCESS;`
`5220`	`5237`	`}`
`5221`	`5238`
`@@ -5265,6 +5282,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,`
`5265`	`5282`	`Py_XDECREF(entry);`
`5266`	`5283`	`Py_XDECREF(inline_state.pushed_locals);`
`5267`	`5284`	`Py_XDECREF(inline_state.temp_symbols);`
	`5285`	`+ Py_XDECREF(inline_state.fast_hidden);`
`5268`	`5286`	`return ERROR;`
`5269`	`5287`	`}`
`5270`	`5288`