gh-124697: avoid unnecessary deduplication in frameobject API functions

Author: iritkatriel

Include/cpython/code.h

@@ -107,6 +107,10 @@ typedef struct {
     struct _PyCoMonitoringData *_co_monitoring; /* Monitoring data */          \
     _PyCode_DEF_UNIQUE_ID()                                                    \
     int _co_firsttraceable;       /* index of first traceable instruction */   \
+    /* 0 = not yet known, 1 = co_localsplusnames are unique, 2 = a name is     \
+       repeated. FrameLocalsProxy learns this from the seen map it builds      \
+       while collecting locals, and skips deduplication when this is 1. */     \
+    char _co_localsplus_dedup;                                                 \
     /* Scratch space for extra data relating to the code object.               \
        Type is a void* to keep the format private in codeobject.c to force     \
        people to go through the proper APIs. */                                \

Include/internal/pycore_code.h

@@ -273,6 +273,11 @@ extern PyObject* _PyCode_GetCellvars(PyCodeObject *);
 extern PyObject* _PyCode_GetFreevars(PyCodeObject *);
 extern PyObject* _PyCode_GetCode(PyCodeObject *);
 
+/* co->_co_localsplus_dedup. Learned by FrameLocalsProxy from its seen map. */
+#define _PYCODE_LOCALSPLUS_DEDUP_UNKNOWN 0
+#define _PYCODE_LOCALSPLUS_DEDUP_UNIQUE  1
+#define _PYCODE_LOCALSPLUS_DEDUP_NEEDED  2
+
 /** API for initializing the line number tables. */
 PyAPI_FUNC(int) _PyCode_InitAddressRange(PyCodeObject* co, PyCodeAddressRange *bounds);
 

Lib/test/test_listcomps.py

@@ -944,6 +944,23 @@ def inner():
             {"snaps": [1, 2], "vals": [2, 2], "consistent": [True, True]},
             ns={"sys": sys}, scopes=["module", "function"])
 
+    def test_frame_locals_dedup_learned_before_cell_is_bound(self):
+        # Reading f_locals before the comprehension cell is bound must not
+        # cache "names are unique": the repeated name is still on the code
+        # object, and dict(**f_locals) inside the comprehension has to work.
+        code = """
+            def outer(x):
+                def inner():
+                    list(sys._getframe().f_locals)
+                    return [dict(**sys._getframe().f_locals)['x'] for x in x]
+                return inner()
+            result = outer([7, 8])
+        """
+        import sys
+        self._check_in_scopes(
+            code, {"result": [7, 8]},
+            ns={"sys": sys}, scopes=["module", "function"])
+
     def test_frame_locals_nested_comp_cell_and_enclosing_free(self):
         # Stress a nested inlined shape where a comp cell and enclosing free
         # share a name; all f_locals views must stay consistent.

Objects/codeobject.c

@@ -561,6 +561,7 @@ init_code(PyCodeObject *co, struct _PyCodeConstructor *con)
     /* not set */
     co->co_weakreflist = NULL;
     co->co_extra = NULL;
+    co->_co_localsplus_dedup = _PYCODE_LOCALSPLUS_DEDUP_UNKNOWN;
     co->_co_cached = NULL;
     co->co_executors = NULL;
 

Objects/frameobject.c

@@ -94,20 +94,62 @@ framelocalsproxy_hasval(_PyInterpreterFrame *frame, PyCodeObject *co, int i)
     return true;
 }
 
+/* *seen is NULL when dedup is known to be unnecessary.
+   Otherwise it maps each localsplus name to Py_True if a value was already
+   kept, or Py_False if the name was only seen on an empty slot. An empty
+   slot must not count as the kept value, but it still proves the name is
+   repeated, so it is recorded while _co_localsplus_dedup is unknown.
+   Returns -1 if the map cannot be created. */
 static int
-framelocalsproxy_is_first_occurrence(PyObject *seen, PyObject *name)
+framelocalsproxy_prepare_dedup(PyCodeObject *co, PyObject **seen)
 {
-    int found = PySet_Contains(seen, name);
-    if (found < 0) {
+    *seen = NULL;
+    if (co->_co_localsplus_dedup == _PYCODE_LOCALSPLUS_DEDUP_UNIQUE) {
+        return 0;
+    }
+    *seen = PyDict_New();
+    if (*seen == NULL) {
         return -1;
     }
-    if (found) {
+    return 0;
+}
+
+/* Returns 1 if this slot's value should be kept, 0 if not, -1 on error.
+   An empty slot is skipped once _co_localsplus_dedup is known. A repeated
+   name stores _PYCODE_LOCALSPLUS_DEDUP_NEEDED. UNIQUE is stored by the
+   caller only after the walk finishes without one. */
+static int
+framelocalsproxy_visit_name(PyCodeObject *co, int i, PyObject *seen,
+                            int has_value)
+{
+    if (seen == NULL) {
+        return has_value;
+    }
+    if (!has_value &&
+        co->_co_localsplus_dedup != _PYCODE_LOCALSPLUS_DEDUP_UNKNOWN) {
         return 0;
     }
-    if (PySet_Add(seen, name) < 0) {
-        return -1;
+    PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
+    PyObject *prev = PyDict_GetItemWithError(seen, name);
+    if (prev == NULL) {
+        if (PyErr_Occurred()) {
+            return -1;
+        }
+        if (PyDict_SetItem(seen, name, has_value ? Py_True : Py_False) < 0) {
+            return -1;
+        }
+        return has_value;
+    }
+    else {
+        co->_co_localsplus_dedup = _PYCODE_LOCALSPLUS_DEDUP_NEEDED;
+        if (!has_value || prev == Py_True) {
+            return 0;
+        }
+        if (PyDict_SetItem(seen, name, Py_True) < 0) {
+            return -1;
+        }
+        return 1;
     }
-    return 1;
 }
 
 static int
@@ -397,27 +439,30 @@ framelocalsproxy_keys(PyObject *self, PyObject *Py_UNUSED(ignored))
         return NULL;
     }
     // An inlined comprehension cell can share a name with a free var.
-    PyObject *seen = PySet_New(NULL);
-    if (seen == NULL) {
+    // Assume dedup may be needed until a walk of seen shows otherwise.
+    PyObject *seen = NULL;
+    if (framelocalsproxy_prepare_dedup(co, &seen) < 0) {
         Py_DECREF(names);
         return NULL;
     }
 
     for (int i = 0; i < co->co_nlocalsplus; i++) {
-        if (framelocalsproxy_hasval(frame->f_frame, co, i)) {
+        int has = framelocalsproxy_hasval(frame->f_frame, co, i);
+        int emit = framelocalsproxy_visit_name(co, i, seen, has);
+        if (emit < 0) {
+            goto error;
+        }
+        if (emit) {
             PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
-            int first = framelocalsproxy_is_first_occurrence(seen, name);
-            if (first < 0) {
+            if (PyList_Append(names, name) < 0) {
                 goto error;
             }
-            if (first) {
-                if (PyList_Append(names, name) < 0) {
-                    goto error;
-                }
-            }
         }
     }
-    Py_DECREF(seen);
+    if (co->_co_localsplus_dedup == _PYCODE_LOCALSPLUS_DEDUP_UNKNOWN) {
+        co->_co_localsplus_dedup = _PYCODE_LOCALSPLUS_DEDUP_UNIQUE;
+    }
+    Py_XDECREF(seen);
 
     // Iterate through the extra locals
     if (frame->f_extra_locals) {
@@ -438,7 +483,7 @@ framelocalsproxy_keys(PyObject *self, PyObject *Py_UNUSED(ignored))
     return names;
 
 error:
-    Py_DECREF(seen);
+    Py_XDECREF(seen);
     Py_DECREF(names);
     return NULL;
 }
@@ -622,30 +667,30 @@ framelocalsproxy_values(PyObject *self, PyObject *Py_UNUSED(ignored))
     if (values == NULL) {
         return NULL;
     }
-    PyObject *seen = PySet_New(NULL);
-    if (seen == NULL) {
+    PyObject *seen = NULL;
+    if (framelocalsproxy_prepare_dedup(co, &seen) < 0) {
         Py_DECREF(values);
         return NULL;
     }
 
     for (int i = 0; i < co->co_nlocalsplus; i++) {
         PyObject *value = framelocalsproxy_getval(frame->f_frame, co, i);
-        if (value) {
-            PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
-            int first = framelocalsproxy_is_first_occurrence(seen, name);
-            if (first == 1) {
-                if (PyList_Append(values, value) < 0) {
-                    Py_DECREF(value);
-                    goto error;
-                }
-            }
+        int has = value != NULL;
+        int emit = framelocalsproxy_visit_name(co, i, seen, has);
+        if (emit < 0) {
+            Py_XDECREF(value);
+            goto error;
+        }
+        if (emit && PyList_Append(values, value) < 0) {
             Py_DECREF(value);
-            if (first < 0) {
-                goto error;
-            }
+            goto error;
         }
+        Py_XDECREF(value);
+    }
+    if (co->_co_localsplus_dedup == _PYCODE_LOCALSPLUS_DEDUP_UNKNOWN) {
+        co->_co_localsplus_dedup = _PYCODE_LOCALSPLUS_DEDUP_UNIQUE;
     }
-    Py_DECREF(seen);
+    Py_XDECREF(seen);
 
     // Iterate through the extra locals
     if (frame->f_extra_locals) {
@@ -663,7 +708,7 @@ framelocalsproxy_values(PyObject *self, PyObject *Py_UNUSED(ignored))
     return values;
 
 error:
-    Py_DECREF(seen);
+    Py_XDECREF(seen);
     Py_DECREF(values);
     return NULL;
 }
@@ -677,36 +722,37 @@ framelocalsproxy_items(PyObject *self, PyObject *Py_UNUSED(ignored))
     if (items == NULL) {
         return NULL;
     }
-    PyObject *seen = PySet_New(NULL);
-    if (seen == NULL) {
+    PyObject *seen = NULL;
+    if (framelocalsproxy_prepare_dedup(co, &seen) < 0) {
         Py_DECREF(items);
         return NULL;
     }
 
     for (int i = 0; i < co->co_nlocalsplus; i++) {
-        PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
         PyObject *value = framelocalsproxy_getval(frame->f_frame, co, i);
-
-        if (value) {
-            int first = framelocalsproxy_is_first_occurrence(seen, name);
-            if (first == 1) {
-                PyObject *pair = _PyTuple_FromPairSteal(Py_NewRef(name), value);
-                if (pair == NULL) {
-                    goto error;
-                }
-                if (_PyList_AppendTakeRef((PyListObject *)items, pair) < 0) {
-                    goto error;
-                }
-            }
-            else {
-                Py_DECREF(value);
-                if (first < 0) {
-                    goto error;
-                }
-            }
+        int has = value != NULL;
+        int emit = framelocalsproxy_visit_name(co, i, seen, has);
+        if (emit < 0) {
+            Py_XDECREF(value);
+            goto error;
         }
+        if (!emit) {
+            Py_XDECREF(value);
+            continue;
+        }
+        PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
+        PyObject *pair = _PyTuple_FromPairSteal(Py_NewRef(name), value);
+        if (pair == NULL) {
+            goto error;
+        }
+        if (_PyList_AppendTakeRef((PyListObject *)items, pair) < 0) {
+            goto error;
+        }
+    }
+    if (co->_co_localsplus_dedup == _PYCODE_LOCALSPLUS_DEDUP_UNKNOWN) {
+        co->_co_localsplus_dedup = _PYCODE_LOCALSPLUS_DEDUP_UNIQUE;
     }
-    Py_DECREF(seen);
+    Py_XDECREF(seen);
     seen = NULL;
 
     // Iterate through the extra locals
@@ -746,24 +792,25 @@ framelocalsproxy_length(PyObject *self)
         size += PyDict_Size(frame->f_extra_locals);
     }
 
-    PyObject *seen = PySet_New(NULL);
-    if (seen == NULL) {
+    PyObject *seen = NULL;
+    if (framelocalsproxy_prepare_dedup(co, &seen) < 0) {
         return -1;
     }
     for (int i = 0; i < co->co_nlocalsplus; i++) {
-        if (framelocalsproxy_hasval(frame->f_frame, co, i)) {
-            PyObject *name = PyTuple_GET_ITEM(co->co_localsplusnames, i);
-            int first = framelocalsproxy_is_first_occurrence(seen, name);
-            if (first < 0) {
-                Py_DECREF(seen);
-                return -1;
-            }
-            else if (first) {
-                size++;
-            }
+        int has = framelocalsproxy_hasval(frame->f_frame, co, i);
+        int emit = framelocalsproxy_visit_name(co, i, seen, has);
+        if (emit < 0) {
+            Py_XDECREF(seen);
+            return -1;
         }
+        if (emit) {
+            size++;
+        }
+    }
+    if (co->_co_localsplus_dedup == _PYCODE_LOCALSPLUS_DEDUP_UNKNOWN) {
+        co->_co_localsplus_dedup = _PYCODE_LOCALSPLUS_DEDUP_UNIQUE;
     }
-    Py_DECREF(seen);
+    Py_XDECREF(seen);
     return size;
 }