gh-99761: Add _PyLong_IsPositiveSingleDigit function to check for single digit integers (#100064)

Author: eendebakpt

Include/internal/pycore_long.h

@@ -110,6 +110,25 @@ PyAPI_FUNC(char*) _PyLong_FormatBytesWriter(
     int base,
     int alternate);
 
+/* Return 1 if the argument is positive single digit int */
+static inline int
+_PyLong_IsPositiveSingleDigit(PyObject* sub) {
+    /*  For a positive single digit int, the value of Py_SIZE(sub) is 0 or 1.
+
+        We perform a fast check using a single comparison by casting from int
+        to uint which casts negative numbers to large positive numbers.
+        For details see Section 14.2 "Bounds Checking" in the Agner Fog
+        optimization manual found at:
+        https://www.agner.org/optimize/optimizing_cpp.pdf
+
+        The function is not affected by -fwrapv, -fno-wrapv and -ftrapv
+        compiler options of GCC and clang
+    */
+    assert(PyLong_CheckExact(sub));
+    Py_ssize_t signed_size = Py_SIZE(sub);
+    return ((size_t)signed_size) <= 1;
+}
+
 #ifdef __cplusplus
 }
 #endif

Python/bytecodes.c

@@ -391,8 +391,7 @@ dummy_func(
             DEOPT_IF(!PyList_CheckExact(list), BINARY_SUBSCR);
 
             // Deopt unless 0 <= sub < PyList_Size(list)
-            Py_ssize_t signed_magnitude = Py_SIZE(sub);
-            DEOPT_IF(((size_t)signed_magnitude) > 1, BINARY_SUBSCR);
+            DEOPT_IF(!_PyLong_IsPositiveSingleDigit(sub), BINARY_SUBSCR);
             assert(((PyLongObject *)_PyLong_GetZero())->ob_digit[0] == 0);
             Py_ssize_t index = ((PyLongObject*)sub)->ob_digit[0];
             DEOPT_IF(index >= PyList_GET_SIZE(list), BINARY_SUBSCR);
@@ -410,8 +409,7 @@ dummy_func(
             DEOPT_IF(!PyTuple_CheckExact(tuple), BINARY_SUBSCR);
 
             // Deopt unless 0 <= sub < PyTuple_Size(list)
-            Py_ssize_t signed_magnitude = Py_SIZE(sub);
-            DEOPT_IF(((size_t)signed_magnitude) > 1, BINARY_SUBSCR);
+            DEOPT_IF(!_PyLong_IsPositiveSingleDigit(sub), BINARY_SUBSCR);
             assert(((PyLongObject *)_PyLong_GetZero())->ob_digit[0] == 0);
             Py_ssize_t index = ((PyLongObject*)sub)->ob_digit[0];
             DEOPT_IF(index >= PyTuple_GET_SIZE(tuple), BINARY_SUBSCR);
@@ -508,7 +506,7 @@ dummy_func(
             DEOPT_IF(!PyList_CheckExact(list), STORE_SUBSCR);
 
             // Ensure nonnegative, zero-or-one-digit ints.
-            DEOPT_IF(((size_t)Py_SIZE(sub)) > 1, STORE_SUBSCR);
+            DEOPT_IF(!_PyLong_IsPositiveSingleDigit(sub), STORE_SUBSCR);
             Py_ssize_t index = ((PyLongObject*)sub)->ob_digit[0];
             // Ensure index < len(list)
             DEOPT_IF(index >= PyList_GET_SIZE(list), STORE_SUBSCR);