From d8802a1d71e40364b89d213dab35a86671d7d6a8 Mon Sep 17 00:00:00 2001
From: Carl Meyer <carl@oddbird.net>
Date: Tue, 7 Mar 2023 10:22:32 -0800
Subject: [PATCH] simplify scope handling

---
 Python/compile.c  | 19 +++++++------------
 Python/symtable.c | 25 +++++++------------------
 2 files changed, 14 insertions(+), 30 deletions(-)

diff --git a/Python/compile.c b/Python/compile.c
index b5d12258cedcfc..7257668c03639e 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -5162,8 +5162,8 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
     assert(c->u->u_fastlocals == NULL);
     if (c->u->u_ste->ste_type != FunctionBlock) {
         // comprehension in non-function scope; for isolation, we'll need to
-        // override names bound in the comprehension to use fast locals, even
-        // though nothing else in this frame will
+        // temporarily override names bound in the comprehension to use fast
+        // locals, even though nothing else in this frame will
         c->u->u_fastlocals = PySet_New(0);
         if (c->u->u_fastlocals == NULL) {
             return ERROR;
@@ -5190,15 +5190,10 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
             }
             assert(PyLong_Check(outv));
             long outsc = (PyLong_AS_LONG(outv) >> SCOPE_OFFSET) & SCOPE_MASK;
-            int outer_global = (outsc == GLOBAL_IMPLICIT || outsc == GLOBAL_EXPLICIT);
-            if (outer_global || ((outsc == CELL || outsc == FREE) && scope == LOCAL)) {
-                // If a name is global in the outer scope but local in the
-                // comprehension scope, we need to keep it global in outer scope
-                // but ensure the comprehension writes to the local, not the
-                // global; this is all the isolation we need.  If it's a cell in
-                // outer scope and a local inside the comprehension, we want the
-                // comprehension to treat it as a local, but we also need to
-                // save/restore the outer cell.
+            if (scope != outsc) {
+                // If a name has different scope inside than outside the
+                // comprehension, we need to temporarily handle it with the
+                // right scope while compiling the comprehension.
                 if (state->temp_symbols == NULL) {
                     state->temp_symbols = PyDict_New();
                     if (state->temp_symbols == NULL) {
@@ -5219,7 +5214,7 @@ push_inlined_comprehension_state(struct compiler *c, location loc,
                 }
                 Py_DECREF(outv);
             }
-            if (!outer_global && (scope == LOCAL || scope == CELL)) {
+            if (outsc == LOCAL || outsc == CELL || outsc == FREE) {
                 // local names bound in comprehension must be isolated from
                 // outer scope; push existing value (which may be NULL if
                 // not defined) on stack
diff --git a/Python/symtable.c b/Python/symtable.c
index df1c0cb24f5144..33b7f335cfe3bc 100644
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -605,23 +605,12 @@ inline_comprehension(PySTEntryObject *ste, PySTEntryObject *comp,
             }
             SET_SCOPE(scopes, k, scope);
         } else {
-            // name already exists in scope
-            PyObject *v_existing_scope = PyDict_GetItemWithError(scopes, k);
-            if (v_existing_scope == NULL) {
-                return 0;
-            }
-            long existing_scope = PyLong_AsLong(v_existing_scope);
-            // if name in comprehension was a cell, promote to cell
-            if (scope == CELL && existing_scope != CELL) {
-                SET_SCOPE(scopes, k, CELL);
-            } else {
-                // free vars in comprehension that are locals in outer scope can
-                // now simply be locals, unless they are free in comp children
-                if ((PyLong_AsLong(existing) & DEF_BOUND) &&
-                     !is_free_in_children(comp, k)) {
-                    if (PySet_Discard(comp_free, 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
+            if ((PyLong_AsLong(existing) & DEF_BOUND) &&
+                    !is_free_in_children(comp, k)) {
+                if (PySet_Discard(comp_free, k) < 0) {
+                    return 0;
                 }
             }
         }
@@ -906,7 +895,7 @@ analyze_block(PySTEntryObject *ste, PyObject *bound, PyObject *free,
         assert(c && PySTEntry_Check(c));
         entry = (PySTEntryObject*)c;
 
-        // we inline comprehensions if inside a function and not a generator
+        // we inline all non-generator-expression comprehensions
         int inline_comp =
             entry->ste_comprehension &&
             !entry->ste_generator;