gh-151289: specialize shrinking int subtraction

Author: KRRT7

Include/internal/pycore_long.h

@@ -116,6 +116,7 @@ PyAPI_DATA(PyObject*) _PyLong_Lshift(PyObject *, int64_t);
 PyAPI_FUNC(_PyStackRef) _PyCompactLong_Add(PyLongObject *left, PyLongObject *right);
 PyAPI_FUNC(_PyStackRef) _PyCompactLong_Multiply(PyLongObject *left, PyLongObject *right);
 PyAPI_FUNC(_PyStackRef) _PyCompactLong_Subtract(PyLongObject *left, PyLongObject *right);
+PyAPI_FUNC(_PyStackRef) _PyLong_SubtractShrink(PyLongObject *left, PyLongObject *right);
 
 // Export for 'binascii' shared extension.
 PyAPI_DATA(unsigned char) _PyLong_DigitValue[256];
@@ -346,6 +347,28 @@ _PyLong_CheckExactAndCompact(PyObject *op)
     return PyLong_CheckExact(op) && _PyLong_IsCompact((const PyLongObject *)op);
 }
 
+static inline bool
+_PyLong_IsShrinkSubtractPair(const PyLongObject *left, const PyLongObject *right)
+{
+#if PyLong_SHIFT == 30
+    if (_PyLong_IsCompact((PyLongObject *)left) ||
+        _PyLong_IsCompact((PyLongObject *)right) ||
+        !_PyLong_SameSign(left, right)) {
+        return false;
+    }
+    Py_ssize_t left_size = _PyLong_DigitCount(left);
+    if (left_size != _PyLong_DigitCount(right) ||
+        left_size < 2 ||
+        left_size > 3) {
+        return false;
+    }
+    return left->long_value.ob_digit[left_size - 1] ==
+           right->long_value.ob_digit[left_size - 1];
+#else
+    return false;
+#endif
+}
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_opcode_metadata.h

@@ -4455,6 +4455,22 @@ def testfunc(n):
         self.assertIn("_POP_TOP_NOP", uops)
         self.assertLessEqual(count_ops(ex, "_POP_TOP"), 2)
 
+    def test_int_subtract_shrink_specialization(self):
+        def testfunc(args):
+            a, b, n = args
+            res = 0
+            for _ in range(n):
+                res = a - b
+            return res
+
+        res, ex = self._run_with_optimizer(
+            testfunc, (10_000_000_001, 10_000_000_000, TIER2_THRESHOLD)
+        )
+        self.assertEqual(res, 1)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertIn("_BINARY_OP_SUBTRACT_INT", uops)
+
     def test_int_mul_op_refcount_elimination(self):
         def testfunc(n):
             c = 1

Include/internal/pycore_uop_ids.h

@@ -1140,6 +1140,25 @@ def test_small_ints(self):
         self.assertIs(i - i, 0)
         self.assertIs(0 * i, 0)
 
+    @support.cpython_only
+    def test_int64_boundary_sub(self):
+        INT64_MAX = (1 << 63) - 1
+        INT64_MIN = -(1 << 63)
+
+        self.assertEqual(INT64_MIN - 0, INT64_MIN)
+        self.assertEqual(INT64_MIN - (-1), INT64_MIN + 1)
+
+        self.assertEqual(INT64_MAX - 0, INT64_MAX)
+        self.assertEqual(INT64_MAX - 1, INT64_MAX - 1)
+
+        self.assertEqual(INT64_MIN - 1, INT64_MIN - 1)
+        self.assertEqual(INT64_MAX - (-1), 1 << 63)
+
+        beyond_max = INT64_MAX + 2
+        beyond_min = INT64_MIN - 2
+        self.assertEqual(beyond_max - 1, INT64_MAX + 1)
+        self.assertEqual(beyond_min - (-1), INT64_MIN - 1)
+
     def test_bit_length(self):
         tiny = 1e-10
         for x in range(-65000, 65000):

Include/internal/pycore_uop_metadata.h

@@ -0,0 +1,2 @@
+Improve performance of binary ``int`` subtraction for shrinking wide exact
+integers.