gh-152666: Avoid reference counting of code objects

Author: brijkapadia

Include/cpython/funcobject.h

@@ -43,6 +43,7 @@ typedef struct {
     PyObject *func_annotations; /* Annotations, a dict or NULL */
     PyObject *func_annotate;    /* Callable to fill the annotations dictionary */
     PyObject *func_typeparams;  /* Tuple of active type variables or NULL */
+    PyObject *func_old_codes;
     vectorcallfunc vectorcall;
     /* Version number for use by specializer.
      * Can set to non-zero when we want to specialize.

Include/internal/pycore_interpframe.h

@@ -132,7 +132,7 @@ _PyFrame_NumSlotsForCodeObject(PyCodeObject *code)
 
 static inline void _PyFrame_Copy(_PyInterpreterFrame *src, _PyInterpreterFrame *dest)
 {
-    dest->f_executable = PyStackRef_MakeHeapSafe(src->f_executable);
+    dest->f_executable = PyStackRef_Borrow(src->f_executable);
     // Don't leave a dangling pointer to the old frame when creating generators
     // and coroutines:
     dest->previous = NULL;
@@ -191,7 +191,7 @@ _PyFrame_Initialize(
 {
     frame->previous = previous;
     frame->f_funcobj = func;
-    frame->f_executable = PyStackRef_FromPyObjectNew(code);
+    frame->f_executable = PyStackRef_FromPyObjectBorrow((PyObject *)code);
     PyFunctionObject *func_obj = (PyFunctionObject *)PyStackRef_AsPyObjectBorrow(func);
     frame->f_builtins = func_obj->func_builtins;
     frame->f_globals = func_obj->func_globals;
@@ -424,7 +424,7 @@ _PyFrame_PushTrampolineUnchecked(PyThreadState *tstate, PyCodeObject *code, int
     assert(tstate->datastack_top < tstate->datastack_limit);
     frame->previous = previous;
     frame->f_funcobj = PyStackRef_None;
-    frame->f_executable = PyStackRef_FromPyObjectNew(code);
+    frame->f_executable = PyStackRef_FromPyObjectBorrow((PyObject *)code);
 #ifdef Py_DEBUG
     frame->f_builtins = NULL;
     frame->f_globals = NULL;

Lib/test/test_capi/test_function.py

@@ -325,6 +325,28 @@ def annofn(arg: int) -> str:
         with self.assertRaises(SystemError):
             _testcapi.function_get_annotations(None)
 
+    def test_function_old_codes(self):
+        def f():
+            pass
+
+        def g():
+            pass
+
+        def h():
+            pass
+
+        old_codes = _testcapi.function_get_old_codes(f)
+        self.assertIsNone(old_codes)
+
+        f.__code__ = g.__code__
+        old_codes = _testcapi.function_get_old_codes(f)
+        self.assertIsInstance(old_codes, list)
+        self.assertEqual(len(old_codes), 1)
+
+        f.__code__ = h.__code__
+        old_codes = _testcapi.function_get_old_codes(f)
+        self.assertEqual(len(old_codes), 2)
+
     # TODO: test PyFunction_New()
     # TODO: test PyFunction_NewWithQualName()
     # TODO: test PyFunction_SetVectorcall()

Modules/_testcapi/function.c

@@ -130,6 +130,19 @@ function_get_annotations(PyObject *self, PyObject *func)
 }
 
 
+static PyObject *
+function_get_old_codes(PyObject *self, PyObject *func)
+{
+    PyFunctionObject *func_o = (PyFunctionObject *) func;
+    PyObject *old_codes = func_o->func_old_codes;
+    if (old_codes == NULL) {
+        Py_RETURN_NONE;
+    }
+
+    return Py_NewRef(old_codes);
+}
+
+
 static PyMethodDef test_methods[] = {
     {"function_get_code", function_get_code, METH_O, NULL},
     {"function_get_globals", function_get_globals, METH_O, NULL},
@@ -141,6 +154,7 @@ static PyMethodDef test_methods[] = {
     {"function_get_closure", function_get_closure, METH_O, NULL},
     {"function_set_closure", function_set_closure, METH_VARARGS, NULL},
     {"function_get_annotations", function_get_annotations, METH_O, NULL},
+    {"function_get_old_codes", function_get_old_codes, METH_O, NULL},
     {NULL},
 };
 

Objects/funcobject.c

@@ -145,6 +145,7 @@ _PyFunction_FromConstructor(PyFrameConstructor *constr)
     op->func_typeparams = NULL;
     op->vectorcall = _PyFunction_Vectorcall;
     op->func_version = FUNC_VERSION_UNSET;
+    op->func_old_codes = NULL;
     // NOTE: functions created via FrameConstructor do not use deferred
     // reference counting because they are typically not part of cycles
     // nor accessed by multiple threads.
@@ -223,6 +224,7 @@ PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname
     op->func_typeparams = NULL;
     op->vectorcall = _PyFunction_Vectorcall;
     op->func_version = FUNC_VERSION_UNSET;
+    op->func_old_codes = NULL;
     if (((code_obj->co_flags & CO_NESTED) == 0) ||
         (code_obj->co_flags & CO_METHOD)) {
         // Use deferred reference counting for top-level functions, but not
@@ -686,6 +688,19 @@ func_set_code(PyObject *self, PyObject *value, void *Py_UNUSED(ignored))
 
     handle_func_event(PyFunction_EVENT_MODIFY_CODE, op, value);
     _PyFunction_ClearVersion(op);
+    if (op->func_old_codes == NULL) {
+        op->func_old_codes = PyList_New(0);
+        if (op->func_old_codes == NULL) {
+            PyErr_NoMemory();
+            return -1;
+        }
+    }
+
+    if (PyList_Append(op->func_old_codes, op->func_code) < 0) {
+            PyErr_NoMemory();
+        return -1;
+    }
+
     Py_XSETREF(op->func_code, Py_NewRef(value));
     return 0;
 }
@@ -1114,6 +1129,7 @@ func_clear(PyObject *self)
     Py_CLEAR(op->func_annotations);
     Py_CLEAR(op->func_annotate);
     Py_CLEAR(op->func_typeparams);
+    Py_CLEAR(op->func_old_codes);
     // Don't Py_CLEAR(op->func_code), since code is always required
     // to be non-NULL. Similarly, name and qualname shouldn't be NULL.
     // However, name and qualname could be str subclasses, so they
@@ -1169,6 +1185,7 @@ func_traverse(PyObject *self, visitproc visit, void *arg)
     Py_VISIT(f->func_annotate);
     Py_VISIT(f->func_typeparams);
     Py_VISIT(f->func_qualname);
+    Py_VISIT(f->func_old_codes);
     return 0;
 }