pythongh-100239: specialize bitwose logical binary ops

iritkatriel · iritkatriel · commit 7a6319b3df89 · 2025-01-16T22:34:43.000Z
diff --git a/Python/specialize.c b/Python/specialize.c
@@ -2414,6 +2414,35 @@ binary_op_fail_kind(int oparg, PyObject *lhs, PyObject *rhs)
 
 /** Binary Op Specialization Extensions */
 
+/* long-long */
+
+static inline int
+is_nonnegative_compactlong(PyObject *v)
+{
+    return PyLong_CheckExact(v) &&
+           (!_PyLong_IsNegative((PyLongObject *)v)) &&
+           _PyLong_IsCompact((PyLongObject *)v);
+}
+
+static int
+nonnegative_compactlongs_guard(PyObject *lhs, PyObject *rhs)
+{
+    return (is_nonnegative_compactlong(lhs) && is_nonnegative_compactlong(rhs));
+}
+
+#define NONNEGATIVE_LONGS_ACTION(NAME, OP) \
+    static PyObject * \
+    (NAME)(PyObject *lhs, PyObject *rhs) \
+    { \
+        Py_ssize_t rhs_val = _PyLong_CompactValue((PyLongObject *)rhs); \
+        Py_ssize_t lhs_val = _PyLong_CompactValue((PyLongObject *)lhs); \
+        return PyLong_FromLong(lhs_val OP rhs_val); \
+    }
+NONNEGATIVE_LONGS_ACTION(nonnegative_compactlongs_or, |)
+NONNEGATIVE_LONGS_ACTION(nonnegative_compactlongs_and, &)
+NONNEGATIVE_LONGS_ACTION(nonnegative_compactlongs_xor, ^)
+#undef NONNEGATIVE_LONGS_ACTION
+
 /* float-long */
 
 static int
@@ -2466,6 +2495,12 @@ LONG_FLOAT_ACTION(compactlong_float_multiply, *)
 LONG_FLOAT_ACTION(compactlong_float_true_div, /)
 #undef LONG_FLOAT_ACTION
 
+static _PyBinaryOpSpecializationDescr nonnegative_compactlongs_specs[NB_OPARG_LAST+1] = {
+    [NB_OR] = {nonnegative_compactlongs_guard, nonnegative_compactlongs_or},
+    [NB_AND] = {nonnegative_compactlongs_guard, nonnegative_compactlongs_and},
+    [NB_XOR] = {nonnegative_compactlongs_guard, nonnegative_compactlongs_xor},
+};
+
 static _PyBinaryOpSpecializationDescr float_compactlong_specs[NB_OPARG_LAST+1] = {
     [NB_ADD] = {float_compactlong_guard, float_compactlong_add},
     [NB_SUBTRACT] = {float_compactlong_guard, float_compactlong_subtract},
@@ -2494,6 +2529,7 @@ binary_op_extended_specialization(PyObject *lhs, PyObject *rhs, int oparg,
 
     LOOKUP_SPEC(compactlong_float_specs, oparg);
     LOOKUP_SPEC(float_compactlong_specs, oparg);
+    LOOKUP_SPEC(nonnegative_compactlongs_specs, oparg);
 #undef LOOKUP_SPEC
     return 0;
 }
