basic array plugin

python · iritkatriel · May 5, 2025 · Apr 24, 2025 · Apr 28, 2025 · Apr 29, 2025
commit 72dfba429a4ddaf44bdcd4b91eec27513954c2b2
diff --git a/Include/internal/pycore_opcode_metadata.h b/Include/internal/pycore_opcode_metadata.h
diff --git a/Include/internal/pycore_uop_metadata.h b/Include/internal/pycore_uop_metadata.h
diff --git a/Modules/arraymodule.c b/Modules/arraymodule.c
@@ -3204,25 +3204,38 @@ do {                                                     \
 } while (0)
 
 static inline int
-array_guard(PyObject *lhs, PyObject *rhs)
+array_subscr_guard(PyObject *lhs, PyObject *rhs)
 {
-     fprintf(stderr, "array_guard\n");
-    return 0;
+    PyObject *exc = PyErr_GetRaisedException();
+    PyObject *module = PyType_GetModuleByDef(Py_TYPE(lhs), &arraymodule);
+    if (module == NULL) {
+        if (!PyErr_Occurred() || PyErr_ExceptionMatches(PyExc_TypeError)) {
+            /* lhs is not an array instance - ignore the TypeError (if any) */
+            PyErr_SetRaisedException(exc);
+            return 0;
+        }
+        else {
+            _PyErr_ChainExceptions1(exc);
+            return -1;
+        }
+    }
+    PyErr_SetRaisedException(exc);
+    return array_Check(lhs, get_array_state(module));
 }
 
 static PyObject *
-array_action(PyObject *lhs, PyObject *rhs)
+array_subscr_action(PyObject *lhs, PyObject *rhs)
 {
-    return NULL;
+    return array_subscr(lhs, rhs);
 }
 
 static int
 array_register_specializations(void)
 {
     _PyBinaryOpSpecializationDescr descr = {
-        .oparg = NB_MULTIPLY,
-        .guard = array_guard,
-        .action = array_action,
+        .oparg = NB_SUBSCR,
+        .guard = array_subscr_guard,
+        .action = array_subscr_action,
     };
     if (_Py_Specialize_AddBinaryOpExtention(&descr) < 0) {
         return -1;

@@ -803,7 +803,8 @@ dummy_func(
             assert(INLINE_CACHE_ENTRIES_BINARY_OP == 5);
             assert(d && d->guard);
             int res = d->guard(left_o, right_o);
-            DEOPT_IF(!res);
+            ERROR_IF(res < 0, error);
+            DEOPT_IF(res == 0);
         }
 
         pure op(_BINARY_OP_EXTEND, (descr/4, left, right -- res)) {
@@ -816,6 +817,7 @@ dummy_func(
 
             PyObject *res_o = d->action(left_o, right_o);
             DECREF_INPUTS();
+            ERROR_IF(res_o == NULL, error);
             res = PyStackRef_FromPyObjectSteal(res_o);
         }
 

diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/specialize.c b/Python/specialize.c
@@ -2534,7 +2534,7 @@ LONG_FLOAT_ACTION(compactlong_float_multiply, *)
 LONG_FLOAT_ACTION(compactlong_float_true_div, /)
 #undef LONG_FLOAT_ACTION
 
-static _PyBinaryOpSpecializationDescr binaryop_extend_descrs[] = {
+static _PyBinaryOpSpecializationDescr binaryop_extend_builtins[] = {
     /* long-long arithmetic */
     {NB_OR, compactlongs_guard, compactlongs_or},
     {NB_AND, compactlongs_guard, compactlongs_and},
@@ -2582,6 +2582,8 @@ binary_op_extended_specialization_from_list(
 {
     for (size_t i = 0; i < size; i++) {
         _PyBinaryOpSpecializationDescr *d = &descrs[i];
+        assert(d != NULL);
+        assert(d->guard != NULL);
         if (d->oparg == oparg && d->guard(lhs, rhs)) {
             *descr = d;
             return 1;
@@ -2594,9 +2596,10 @@ 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_descrs,
-            sizeof(binaryop_extend_descrs)/sizeof(_PyBinaryOpSpecializationDescr),
+            binaryop_extend_builtins,
+            sizeof(binaryop_extend_builtins)/sizeof(descr_type),
             lhs, rhs, oparg, descr))
     {
         return 1;
@@ -2605,7 +2608,7 @@ binary_op_extended_specialization(PyObject *lhs, PyObject *rhs, int oparg,
     PyThreadState *tstate = PyThreadState_Get();
     _Py_c_array_t *extensions = &tstate->interp->binop_specializer_extentions;
     if (binary_op_extended_specialization_from_list(
-            (_PyBinaryOpSpecializationDescr *)extensions->array,
+            (descr_type*)extensions->array,
             tstate->interp->num_binop_specializer_extentions,
             lhs, rhs, oparg, descr))
     {