Add space for set-awaiter methods on types without changing existing …

…structs

Summary:
This is part of re-creating the `cr_awaiter` feature backed out in D59866335. The issue was changing the `PyAsyncMethods` struct leads to ABI breakage. In this diff we just make the change to store the extra field we wanted to add in data "outside" of the `PyTypeObject`.

I didn't come up with this; it's mostly a reimplementation for how things were already done in Cinder. We were trying to get away from this while upstreaming this to 3.12 but the patch was rejected.

Reviewed By: alexmalyshev

Differential Revision: D60097673

fbshipit-source-id: 19cabdfec0cfbab35f526494e63a75501b1719d9

Include/cpython/object.h

@@ -132,6 +132,18 @@ typedef struct {
     sendfunc am_send;
 } PyAsyncMethods;
 
+typedef void (*setawaiterfunc)(PyObject *receiver, PyObject *awaiter);
+
+typedef struct {
+    PyAsyncMethods ame_async_methods;
+    setawaiterfunc ame_setawaiter;
+} Ci_AsyncMethodsWithExtra;
+
+#define Ci_HeapType_AM_EXTRA(etype) \
+    ((Ci_AsyncMethodsWithExtra *)(((char *)etype) +  \
+      Py_TYPE(etype)->tp_basicsize + \
+      Py_SIZE(etype) * sizeof(PyMemberDef)))
+
 typedef struct {
      getbufferproc bf_getbuffer;
      releasebufferproc bf_releasebuffer;

Include/object.h

@@ -527,6 +527,9 @@ given type object has a specified feature.
 /* Objects behave like an unbound method */
 #define Py_TPFLAGS_METHOD_DESCRIPTOR (1UL << 17)
 
+/* Set if type's tp_as_async slot points to Ci_AsyncMethodsWithExtra */
+#define Ci_TPFLAGS_HAVE_AM_EXTRA (1UL << 18)
+
 /* Object has up-to-date type attribute cache */
 #define Py_TPFLAGS_VALID_VERSION_TAG  (1UL << 19)
 

Objects/typeobject.c

@@ -3318,21 +3318,45 @@ type_new_alloc(type_new_ctx *ctx)
     PyTypeObject *metatype = ctx->metatype;
     PyTypeObject *type;
 
+    /* If the base class has Ci_AsyncMethodsWithExtra, we allocate space at the
+     * end of this type so it can also have it. The extra slots in
+     * Ci_AsyncMethodsWithExtra aren't added to slotdefs and don't have managed
+     * counterparts, so this has no implications on slotptr() or the relative
+     * order of the various *Methods members of PyHeapTypeObject. */
+    int have_am_extra = PyType_HasFeature(ctx->base, Ci_TPFLAGS_HAVE_AM_EXTRA);
+
+    /* Allocate the type object */
+    Py_ssize_t extra_bytes =
+        (have_am_extra ? sizeof(Ci_AsyncMethodsWithExtra) : 0);
+    Py_ssize_t extra_slots =
+        (extra_bytes + sizeof(PyMemberDef) - 1) / sizeof(PyMemberDef);
+
     // Allocate the type object
-    type = (PyTypeObject *)metatype->tp_alloc(metatype, ctx->nslot);
+    type = (PyTypeObject *)metatype->tp_alloc(metatype, ctx->nslot + extra_slots);
     if (type == NULL) {
         return NULL;
     }
+    Py_SET_SIZE(type, Py_SIZE(type) - extra_slots);
+
     PyHeapTypeObject *et = (PyHeapTypeObject *)type;
 
     // Initialize tp_flags.
     // All heap types need GC, since we can create a reference cycle by storing
     // an instance on one of its parents.
     type->tp_flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HEAPTYPE |
-                      Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC);
+                      Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC |
+                      (ctx->base->tp_flags & Ci_TPFLAGS_HAVE_AM_EXTRA));
 
     // Initialize essential fields
-    type->tp_as_async = &et->as_async;
+    if (have_am_extra) {
+        type->tp_as_async = (PyAsyncMethods *)Ci_HeapType_AM_EXTRA(type);
+        /* Only ame_setawaiter is inherited and it has no managed counterpart,
+         * so it's special-cased here. */
+        ((Ci_AsyncMethodsWithExtra *)type->tp_as_async)->ame_setawaiter =
+            ((Ci_AsyncMethodsWithExtra *)ctx->base->tp_as_async)->ame_setawaiter;
+    } else {
+        type->tp_as_async = &et->as_async;
+    }
     type->tp_as_number = &et->as_number;
     type->tp_as_sequence = &et->as_sequence;
     type->tp_as_mapping = &et->as_mapping;