Fix reference count errors with PyCode_* functions

Summary:
These functions return strong references on Python 3.11 and later.  The
compatibility shims we had in cinderx/Common/code.h were wrong, they returned
borrowed references, and we wrote code expecting their results to be borrowed.

Switch the compatibility shims to use strong references and fix callsites to
properly decref them as needed.  Callsites in shadowcode and the 3.10
interpreter can just be inlined and use borrowed references. That code plus the
shims will go away once 3.10 support goes away (hopefully soon).

The issue was pointed out to me by Claude.  The solution it generated was a
mess, so I handwrote this instead.

Reviewed By: kddnewton, jbower-fb

Differential Revision: D94730099

fbshipit-source-id: 79e6ec81f63314500cdd8bded6792fc95e760250

Author: alexmalyshev

cinderx/Common/code.cpp

@@ -52,9 +52,10 @@ PyObject* getVarnameTuple(BorrowedRef<PyCodeObject> code, int* idx) {
   }
 
   *idx -= code->co_nlocals;
-  auto ncellvars = PyTuple_GET_SIZE(PyCode_GetCellvars(code));
+  auto cellvars = Ref<>::steal(PyCode_GetCellvars(code));
+  auto ncellvars = PyTuple_GET_SIZE(cellvars.get());
   if (*idx < ncellvars) {
-    return PyCode_GetCellvars(code);
+    return cellvars.release();
   }
 
   *idx -= ncellvars;
@@ -65,14 +66,14 @@ PyObject* getVarname(BorrowedRef<PyCodeObject> code, int idx) {
 #if PY_VERSION_HEX >= 0x030C0000
   return PyTuple_GET_ITEM(code->co_localsplusnames, idx);
 #else
-  PyObject* tuple = getVarnameTuple(code, &idx);
-  return PyTuple_GET_ITEM(tuple, idx);
+  auto tuple = Ref<>::steal(getVarnameTuple(code, &idx));
+  return PyTuple_GET_ITEM(tuple.get(), idx);
 #endif
 }
 
 uint32_t hashBytecode(BorrowedRef<PyCodeObject> code) {
   uint32_t crc = crc32(0, nullptr, 0);
-  PyObject* bc = PyCode_GetCode(code);
+  auto bc = Ref<>::steal(PyCode_GetCode(code));
   if (!PyBytes_Check(bc)) {
     return crc;
   }
@@ -111,11 +112,11 @@ _Py_CODEUNIT* codeUnit(PyCodeObject* code) {
 #if PY_VERSION_HEX >= 0x030C0000
   return _PyCode_CODE(code);
 #else
-  PyObject* bytes_obj = PyCode_GetCode(code);
+  auto bytes_obj = Ref<>::steal(PyCode_GetCode(code));
   JIT_DCHECK(
       PyBytes_CheckExact(bytes_obj),
       "Code object must have its instructions stored as a byte string");
-  return (_Py_CODEUNIT*)PyBytes_AS_STRING(PyCode_GetCode(code));
+  return (_Py_CODEUNIT*)PyBytes_AS_STRING(bytes_obj.get());
 #endif
 }
 
@@ -125,7 +126,8 @@ size_t countIndices(PyCodeObject* code) {
   // which we don't need just to determine the number of indices.
   return _PyCode_NBYTES(code) / sizeof(_Py_CODEUNIT);
 #else
-  return PyBytes_GET_SIZE(PyCode_GetCode(code)) / sizeof(_Py_CODEUNIT);
+  auto bytes_obj = Ref<>::steal(PyCode_GetCode(code));
+  return PyBytes_GET_SIZE(bytes_obj.get()) / sizeof(_Py_CODEUNIT);
 #endif
 }
 
@@ -268,15 +270,17 @@ int numCellvars(PyCodeObject* code) {
 #if PY_VERSION_HEX >= 0x030B0000
   return code->co_ncellvars;
 #else
-  return PyTuple_GET_SIZE(PyCode_GetCellvars(code));
+  auto cellvars = Ref<>::steal(PyCode_GetCellvars(code));
+  return PyTuple_GET_SIZE(cellvars.get());
 #endif
 }
 
 int numFreevars(PyCodeObject* code) {
 #if PY_VERSION_HEX >= 0x030B0000
   return code->co_nfreevars;
 #else
-  return PyTuple_GET_SIZE(PyCode_GetFreevars(code));
+  auto freevars = Ref<>::steal(PyCode_GetFreevars(code));
+  return PyTuple_GET_SIZE(freevars.get());
 #endif
 }
 

cinderx/Common/code.h

@@ -21,19 +21,19 @@ extern "C" {
 #if PY_VERSION_HEX < 0x030B0000
 
 static inline PyObject* PyCode_GetCode(PyCodeObject* code) {
-  return code->co_code;
+  return Py_NewRef(code->co_code);
 }
 
 static inline PyObject* PyCode_GetVarnames(PyCodeObject* code) {
-  return code->co_varnames;
+  return Py_NewRef(code->co_varnames);
 }
 
 static inline PyObject* PyCode_GetCellvars(PyCodeObject* code) {
-  return code->co_cellvars;
+  return Py_NewRef(code->co_cellvars);
 }
 
 static inline PyObject* PyCode_GetFreevars(PyCodeObject* code) {
-  return code->co_freevars;
+  return Py_NewRef(code->co_freevars);
 }
 
 static inline PyCodeObject* PyUnstable_Code_New(
@@ -163,8 +163,8 @@ std::string codeFullname(
 std::string funcFullname(BorrowedRef<PyFunctionObject> func);
 
 // Given a code object and an index into f_localsplus, compute which of
-// code->co_varnames, code->cellvars, or code->freevars contains the name of
-// the variable. Return that tuple and adjust idx as needed.
+// code->co_varnames, code->cellvars, or code->freevars contains the name of the
+// variable. Return a new reference to that tuple and adjust idx as needed.
 PyObject* getVarnameTuple(BorrowedRef<PyCodeObject> code, int* idx);
 
 // Similar to getVarnameTuple, but return the name itself rather than the

cinderx/Immortalize/immortalize.cpp

@@ -44,10 +44,10 @@ bool immortalize(PyObject* obj) {
 #if PY_VERSION_HEX < 0x030B0000
     // In 3.11 these changed to have the bytes embedded in the code object and
     // the names in a unified tuple
-    IMMORTALIZE(PyCode_GetCode(code));
-    IMMORTALIZE(PyCode_GetVarnames(code));
-    IMMORTALIZE(PyCode_GetFreevars(code));
-    IMMORTALIZE(PyCode_GetCellvars(code));
+    IMMORTALIZE(code->co_code);
+    IMMORTALIZE(code->co_varnames);
+    IMMORTALIZE(code->co_freevars);
+    IMMORTALIZE(code->co_cellvars);
 #else
     IMMORTALIZE(code->co_localspluskinds);
     IMMORTALIZE(code->co_localsplusnames);

cinderx/Interpreter/3.10/interpreter.c

@@ -310,11 +310,11 @@ Ci_EvalFrame(PyThreadState* tstate, PyFrameObject* f, int throwflag) {
   consts = co->co_consts;
   fastlocals = f->f_localsplus;
   freevars = f->f_localsplus + co->co_nlocals;
-  assert(PyBytes_Check(PyCode_GetCode(co)));
-  assert(PyBytes_GET_SIZE(PyCode_GetCode(co)) <= INT_MAX);
-  assert(PyBytes_GET_SIZE(PyCode_GetCode(co)) % sizeof(_Py_CODEUNIT) == 0);
+  assert(PyBytes_Check(co->co_code));
+  assert(PyBytes_GET_SIZE(co->co_code) <= INT_MAX);
+  assert(PyBytes_GET_SIZE(co->co_code) % sizeof(_Py_CODEUNIT) == 0);
   assert(_Py_IS_ALIGNED(
-      PyBytes_AS_STRING(PyCode_GetCode(co)), sizeof(_Py_CODEUNIT)));
+      PyBytes_AS_STRING(co->co_code), sizeof(_Py_CODEUNIT)));
 
   /* facebook begin t39538061 */
   shadow.code = co;
@@ -504,7 +504,7 @@ Ci_EvalFrame(PyThreadState* tstate, PyFrameObject* f, int throwflag) {
               tstate,
               PyExc_UnboundLocalError,
               UNBOUNDLOCAL_ERROR_MSG,
-              PyTuple_GetItem(PyCode_GetVarnames(co), oparg));
+              PyTuple_GetItem(co->co_varnames, oparg));
           goto error;
         }
         Py_INCREF(value);
@@ -1636,10 +1636,10 @@ Ci_EvalFrame(PyThreadState* tstate, PyFrameObject* f, int throwflag) {
         PyObject *name, *value, *locals = f->f_locals;
         Py_ssize_t idx;
         assert(locals);
-        assert(oparg >= PyTuple_GET_SIZE(PyCode_GetCellvars(co)));
-        idx = oparg - PyTuple_GET_SIZE(PyCode_GetCellvars(co));
-        assert(idx >= 0 && idx < PyTuple_GET_SIZE(PyCode_GetFreevars(co)));
-        name = PyTuple_GET_ITEM(PyCode_GetFreevars(co), idx);
+        assert(oparg >= PyTuple_GET_SIZE(co->co_cellvars));
+        idx = oparg - PyTuple_GET_SIZE(co->co_cellvars);
+        assert(idx >= 0 && idx < PyTuple_GET_SIZE(co->co_freevars));
+        name = PyTuple_GET_ITEM(co->co_freevars, idx);
         if (PyDict_CheckExact(locals)) {
           value = PyDict_GetItemWithError(locals, name);
           if (value != NULL) {
@@ -5020,8 +5020,8 @@ _Ci_CheckArgs(PyThreadState* tstate, PyFrameObject* f, PyCodeObject* co) {
             "%U expected '%s' for argument %U, got '%s'",
             co->co_name,
             type->tp_name,
-            idx < 0 ? PyTuple_GetItem(PyCode_GetCellvars(co), -(idx + 1))
-                    : PyTuple_GetItem(PyCode_GetVarnames(co), idx),
+            idx < 0 ? PyTuple_GetItem(co->co_cellvars, -(idx + 1))
+                    : PyTuple_GetItem(co->co_varnames, idx),
             Py_TYPE(val)->tp_name);
         Py_DECREF(type);
         return -1;
@@ -5070,8 +5070,8 @@ _Ci_CheckArgs(PyThreadState* tstate, PyFrameObject* f, PyCodeObject* co) {
           "%U expected '%s' for argument %U, got '%s'",
           co->co_name,
           check->tai_type->tp_name,
-          idx < 0 ? PyTuple_GetItem(PyCode_GetCellvars(co), -(idx + 1))
-                  : PyTuple_GetItem(PyCode_GetVarnames(co), idx),
+          idx < 0 ? PyTuple_GetItem(co->co_cellvars, -(idx + 1))
+                  : PyTuple_GetItem(co->co_varnames, idx),
           Py_TYPE(val)->tp_name);
       return -1;
     }

cinderx/Jit/hir/hir.cpp

@@ -1281,7 +1281,7 @@ const Environment::ReferenceSet& Environment::references() const {
 
 bool usesRuntimeFunc([[maybe_unused]] BorrowedRef<PyCodeObject> code) {
 #if PY_VERSION_HEX < 0x030C0000
-  return PyTuple_GET_SIZE(PyCode_GetFreevars(code)) > 0;
+  return numFreevars(code) > 0;
 #else
   // In 3.12+ we always need the runtime function because we use it to
   // initialize the _PyInterpreterFrame object.

cinderx/Jit/hir/inliner.cpp

@@ -154,12 +154,10 @@ bool canInline(Function& caller, AbstractCall* call_instr) {
     }
   }
 #else
-  Py_ssize_t ncellvars = PyTuple_GET_SIZE(PyCode_GetCellvars(code));
-  if (ncellvars > 0) {
+  if (numCellvars(code) > 0) {
     return fail(InlineFailureType::kHasCellvars);
   }
-  Py_ssize_t nfreevars = PyTuple_GET_SIZE(PyCode_GetFreevars(code));
-  if (nfreevars > 0) {
+  if (numFreevars(code) > 0) {
     return fail(InlineFailureType::kHasFreevars);
   }
 #endif

cinderx/Jit/hir/printer.cpp

@@ -231,7 +231,7 @@ format_varname(const Function* func, const Instr& instr, int idx) {
     return fmt::format("{}", idx);
   }
 
-  auto names = getVarnameTuple(code, &idx);
+  auto names = Ref<>::steal(getVarnameTuple(code, &idx));
   return format_name_impl(idx, names);
 }
 

cinderx/RuntimeTests/deopt_test.cpp

@@ -347,8 +347,9 @@ def test(x, y):
     const int jump_index = 42;
     const int pop_instr_offset = 2;
 #endif
+    auto code_bytes = Ref<>::steal(PyCode_GetCode(code));
     ASSERT_EQ(
-        PyBytes_AS_STRING(PyCode_GetCode(code))[jump_index + pop_instr_offset],
+        PyBytes_AS_STRING(code_bytes.get())[jump_index + pop_instr_offset],
         (char)POP_JUMP_IF_ZERO);
 
     DeoptMetadata dm;

cinderx/Shadowcode/shadowcode.c

@@ -619,7 +619,7 @@ int _PyShadow_GetOriginalOparg(
     _PyShadow_EvalState* state,
     const _Py_CODEUNIT* next_instr) {
   _Py_CODEUNIT* rawcode =
-      (_Py_CODEUNIT*)PyBytes_AS_STRING(PyCode_GetCode(state->code));
+      (_Py_CODEUNIT*)PyBytes_AS_STRING(state->code->co_code);
   _Py_CODEUNIT* instr = &rawcode[next_instr - *state->first_instr];
   instr--; /* we point to the next instruction, we want the current one */
   const int existing_size = opsize(instr, rawcode);
@@ -2329,7 +2329,7 @@ PyObject* _PyShadow_GetInlineCacheStats(PyObject* self) {
 #endif
 
 int _PyShadow_InitCache(PyCodeObject* co) {
-  char* buffer = PyBytes_AS_STRING(PyCode_GetCode(co));
+  char* buffer = PyBytes_AS_STRING(co->co_code);
 
   /* names is a bitmask of seen names, if we have a limited number */
   size_t names = 0, funcs = 0;
@@ -2356,7 +2356,7 @@ int _PyShadow_InitCache(PyCodeObject* co) {
   /* Scan the byte code for all LOAD_GLOBALs and pre-allocate enough space
    * for all of them */
   _Py_CODEUNIT* instr = (_Py_CODEUNIT*)buffer;
-  _Py_CODEUNIT* end = (_Py_CODEUNIT*)(buffer + Py_SIZE(PyCode_GetCode(co)));
+  _Py_CODEUNIT* end = (_Py_CODEUNIT*)(buffer + Py_SIZE(co->co_code));
   while (instr < end) {
     unsigned char opcode = _Py_OPCODE(*instr);
     int oparg = _Py_OPARG(*instr);
@@ -2407,13 +2407,13 @@ int _PyShadow_InitCache(PyCodeObject* co) {
   }
 
   _PyShadowCode* shadow;
-  shadow = PyMem_Malloc(sizeof(_PyShadowCode) + Py_SIZE(PyCode_GetCode(co)));
+  shadow = PyMem_Malloc(sizeof(_PyShadowCode) + Py_SIZE(co->co_code));
   if (shadow == NULL) {
     return -1;
   }
 
   shadow->update_count = 0;
-  shadow->len = Py_SIZE(PyCode_GetCode(co));
+  shadow->len = Py_SIZE(co->co_code);
   memcpy(shadow->code, buffer, shadow->len);
 
   if (glob_count) {