method on type

Author: iritkatriel

Include/cpython/object.h

@@ -143,6 +143,11 @@ typedef struct {
  * backwards-compatibility */
 typedef Py_ssize_t printfunc;
 
+/* Specialize a binary by setting the descriptor pointer */
+struct _BinaryOpSpecializationDescr;
+typedef int(*binop_specialize_func)(PyObject *v, PyObject *w, int oparg,
+                                    struct _BinaryOpSpecializationDescr **descr);
+
 // If this structure is modified, Doc/includes/typestruct.h should be updated
 // as well.
 struct _typeobject {
@@ -233,6 +238,9 @@ struct _typeobject {
     /* bitset of which type-watchers care about this type */
     unsigned char tp_watched;
 
+    /* callback that may specialize BINARY_OP */
+    binop_specialize_func tp_binop_specialize;
+
     /* Number of tp_version_tag values used.
      * Set to _Py_ATTR_CACHE_UNUSED if the attribute cache is
      * disabled for this type (e.g. due to custom MRO entries).

Include/internal/pycore_code.h

@@ -485,14 +485,12 @@ adaptive_counter_backoff(_Py_BackoffCounter counter) {
 typedef int (*binaryopguardfunc)(PyObject *lhs, PyObject *rhs);
 typedef PyObject *(*binaryopactionfunc)(PyObject *lhs, PyObject *rhs);
 
-typedef struct {
+typedef struct _BinaryOpSpecializationDescr {
     int oparg;
     binaryopguardfunc guard;
     binaryopactionfunc action;
 } _PyBinaryOpSpecializationDescr;
 
-PyAPI_DATA(int) _Py_Specialize_AddBinaryOpExtention(_PyBinaryOpSpecializationDescr* descr);
-
 /* Comparison bit masks. */
 
 /* Note this evaluates its arguments twice each */

Include/typeslots.h

@@ -93,4 +93,5 @@
 #if !defined(Py_LIMITED_API) || Py_LIMITED_API+0 >= 0x030E0000
 /* New in 3.14 */
 #define Py_tp_token 83
+#define Py_tp_binop_specialize 84
 #endif

Lib/test/test_sys.py

@@ -1776,7 +1776,7 @@ def delx(self): del self.__x
         check((1,2,3), vsize('') + self.P + 3*self.P)
         # type
         # static type: PyTypeObject
-        fmt = 'P2nPI13Pl4Pn9Pn12PIPc'
+        fmt = 'P2nPI13Pl4Pn9Pn12PI3Pc'
         s = vsize(fmt)
         check(int, s)
         typeid = 'n' if support.Py_GIL_DISABLED else ''

Modules/arraymodule.c

@@ -850,6 +850,10 @@ array_richcompare(PyObject *v, PyObject *w, int op)
     return res;
 }
 
+static int
+array_binop_specialize(PyObject *v, PyObject *w, int oparg,
+                       _PyBinaryOpSpecializationDescr **descr);
+
 static Py_ssize_t
 array_length(PyObject *op)
 {
@@ -2966,6 +2970,7 @@ static PyType_Slot array_slots[] = {
     {Py_tp_new, array_new},
     {Py_tp_traverse, array_tp_traverse},
     {Py_tp_token, Py_TP_USE_SPEC},
+    {Py_tp_binop_specialize, array_binop_specialize},
 
     /* as sequence */
     {Py_sq_length, array_length},
@@ -2998,6 +3003,61 @@ static PyType_Spec array_spec = {
     .slots = array_slots,
 };
 
+static inline int
+array_subscr_guard(PyObject *lhs, PyObject *rhs)
+{
+    PyObject *exc = PyErr_GetRaisedException();
+    int ret = PyType_GetBaseByToken(Py_TYPE(lhs), &array_spec, NULL);
+    if (ret < 0) {
+        if (PyErr_ExceptionMatches(PyExc_TypeError)) {
+            PyErr_Clear();
+            ret = 0;
+        }
+    }
+    _PyErr_ChainExceptions1(exc);
+    return ret;
+}
+
+static PyObject *
+array_subscr_action(PyObject *lhs, PyObject *rhs)
+{
+    return array_subscr(lhs, rhs);
+}
+
+static int
+array_binop_specialize(PyObject *v, PyObject *w, int oparg,
+                       _PyBinaryOpSpecializationDescr **descr)
+{
+    array_state *state = find_array_state_by_type(Py_TYPE(v));
+
+    if (!array_Check(v, state)) {
+        return 0;
+    }
+
+    *descr = NULL;
+    switch(oparg) {
+        case NB_SUBSCR:
+            if (array_subscr_guard(v, w)) {
+                *descr = (_PyBinaryOpSpecializationDescr*)PyMem_Malloc(
+                             sizeof(_PyBinaryOpSpecializationDescr));
+                if (*descr == NULL) {
+                    PyErr_NoMemory();
+                    return -1;
+                }
+                **descr = (_PyBinaryOpSpecializationDescr) {
+                    .oparg = oparg,
+                    .guard = array_subscr_guard,
+                    .action = array_subscr_action,
+                };
+                return 1;
+            }
+            break;
+    }
+
+    return 0;
+}
+
+
 /*********************** Array Iterator **************************/
 
 /*[clinic input]
@@ -3204,41 +3264,6 @@ do {                                                     \
     state->str_ ## string = tmp;                         \
 } while (0)
 
-static inline int
-array_subscr_guard(PyObject *lhs, PyObject *rhs)
-{
-    PyObject *exc = PyErr_GetRaisedException();
-    int ret = PyType_GetBaseByToken(Py_TYPE(lhs), &array_spec, NULL);
-    if (ret < 0) {
-        if (PyErr_ExceptionMatches(PyExc_TypeError)) {
-            PyErr_Clear();
-            ret = 0;
-        }
-    }
-    _PyErr_ChainExceptions1(exc);
-    return ret;
-}
-
-static PyObject *
-array_subscr_action(PyObject *lhs, PyObject *rhs)
-{
-    return array_subscr(lhs, rhs);
-}
-
-static int
-array_register_specializations(void)
-{
-    _PyBinaryOpSpecializationDescr descr = {
-        .oparg = NB_SUBSCR,
-        .guard = array_subscr_guard,
-        .action = array_subscr_action,
-    };
-    if (_Py_Specialize_AddBinaryOpExtention(&descr) < 0) {
-        return -1;
-    }
-    return 0;
-}
-
 static int
 array_modexec(PyObject *m)
 {
@@ -3278,10 +3303,6 @@ array_modexec(PyObject *m)
     }
     Py_DECREF(res);
 
-    if (array_register_specializations() < 0) {
-        return -1;
-    }
-
     if (PyModule_AddType(m, state->ArrayType) < 0) {
         return -1;
     }

Objects/typeslots.inc

@@ -801,7 +801,8 @@ dummy_func(
             PyObject *right_o = PyStackRef_AsPyObjectBorrow(right);
             _PyBinaryOpSpecializationDescr *d = (_PyBinaryOpSpecializationDescr*)descr;
             assert(INLINE_CACHE_ENTRIES_BINARY_OP == 5);
-            assert(d && d->guard);
+            assert(d);
+            assert(d->guard);
             int res = d->guard(left_o, right_o);
             ERROR_IF(res < 0, error);
             DEOPT_IF(res == 0);

Python/bytecodes.c

@@ -2556,63 +2556,29 @@ static _PyBinaryOpSpecializationDescr binaryop_extend_builtins[] = {
     {NB_MULTIPLY, compactlong_float_guard, compactlong_float_multiply},
 };
 
-int
-_Py_Specialize_AddBinaryOpExtention(_PyBinaryOpSpecializationDescr* descr)
-{
-    PyThreadState *tstate = PyThreadState_Get();
-    _Py_c_array_t *extensions = &tstate->interp->binop_specializer_extentions;
-    Py_ssize_t idx = tstate->interp->num_binop_specializer_extentions;
-    if (idx == 0) {
-        _Py_CArray_Init(extensions, sizeof(_PyBinaryOpSpecializationDescr), 10);
-    }
-    if (_Py_CArray_EnsureCapacity(extensions, idx) < 0) {
-        return -1;
-    }
-    _PyBinaryOpSpecializationDescr* descrs = (_PyBinaryOpSpecializationDescr*)extensions->array;
-    descrs[idx] = *descr;
-    tstate->interp->num_binop_specializer_extentions++;
-    return 0;
-}
-
 static int
-binary_op_extended_specialization_from_list(
-        _PyBinaryOpSpecializationDescr *descrs, size_t size,
-        PyObject *lhs, PyObject *rhs, int oparg,
-        _PyBinaryOpSpecializationDescr **descr)
+binary_op_extended_specialization(PyObject *lhs, PyObject *rhs, int oparg,
+                                  _PyBinaryOpSpecializationDescr **descr)
 {
+    typedef _PyBinaryOpSpecializationDescr descr_type;
+    size_t size = sizeof(binaryop_extend_builtins)/sizeof(descr_type);
     for (size_t i = 0; i < size; i++) {
-        _PyBinaryOpSpecializationDescr *d = &descrs[i];
+        _PyBinaryOpSpecializationDescr *d = &binaryop_extend_builtins[i];
         assert(d != NULL);
         assert(d->guard != NULL);
         if (d->oparg == oparg && d->guard(lhs, rhs)) {
             *descr = d;
             return 1;
         }
     }
-    return 0;
-}
 
-static int
-binary_op_extended_specialization(PyObject *lhs, PyObject *rhs, int oparg,
-                                  _PyBinaryOpSpecializationDescr **descr)
-{
-    typedef _PyBinaryOpSpecializationDescr descr_type;
-    if (binary_op_extended_specialization_from_list(
-            binaryop_extend_builtins,
-            sizeof(binaryop_extend_builtins)/sizeof(descr_type),
-            lhs, rhs, oparg, descr))
-    {
-        return 1;
-    }
+    if (Py_TYPE(lhs)->tp_binop_specialize != NULL) {
+        int ret = Py_TYPE(lhs)->tp_binop_specialize(lhs, rhs, oparg, descr);
+        if (ret < 0) {
+            return -1;
+        }
 
-    PyThreadState *tstate = PyThreadState_Get();
-    _Py_c_array_t *extensions = &tstate->interp->binop_specializer_extentions;
-    if (binary_op_extended_specialization_from_list(
-            (descr_type*)extensions->array,
-            tstate->interp->num_binop_specializer_extentions,
-            lhs, rhs, oparg, descr))
-    {
-        return 1;
+        return ret;
     }
     return 0;
 }