GH-101291: Add low level, unstable API for pylong (GH-101685)

Co-authored-by: Petr Viktorin <[email protected]>

Author: markshannon

Doc/c-api/long.rst

@@ -322,3 +322,27 @@ distinguished from a number.  Use :c:func:`PyErr_Occurred` to disambiguate.
    with :c:func:`PyLong_FromVoidPtr`.
 
    Returns ``NULL`` on error.  Use :c:func:`PyErr_Occurred` to disambiguate.
+
+
+.. c:function:: int PyUnstable_Long_IsCompact(const PyLongObject* op)
+
+   Return 1 if *op* is compact, 0 otherwise.
+
+   This function makes it possible for performance-critical code to implement
+   a “fast path” for small integers. For compact values use
+   :c:func:`PyUnstable_Long_CompactValue`; for others fall back to a
+   :c:func:`PyLong_As* <PyLong_AsSize_t>` function or
+   :c:func:`calling <PyObject_CallMethod>` :meth:`int.to_bytes`.
+
+   The speedup is expected to be negligible for most users.
+
+   Exactly what values are considered compact is an implementation detail
+   and is subject to change.
+
+.. c:function:: Py_ssize_t PyUnstable_Long_CompactValue(const PyLongObject* op)
+
+   If *op* is compact, as determined by :c:func:`PyUnstable_Long_IsCompact`,
+   return its value.
+
+   Otherwise, the return value is undefined.
+

Include/cpython/longintrepr.h

@@ -98,6 +98,32 @@ PyAPI_FUNC(PyLongObject *)
 _PyLong_FromDigits(int negative, Py_ssize_t digit_count, digit *digits);
 
 
+/* Inline some internals for speed. These should be in pycore_long.h
+ * if user code didn't need them inlined. */
+
+#define _PyLong_SIGN_MASK 3
+#define _PyLong_NON_SIZE_BITS 3
+
+static inline int
+_PyLong_IsCompact(const PyLongObject* op) {
+    assert(PyLong_Check(op));
+    return op->long_value.lv_tag < (2 << _PyLong_NON_SIZE_BITS);
+}
+
+#define PyUnstable_Long_IsCompact _PyLong_IsCompact
+
+static inline Py_ssize_t
+_PyLong_CompactValue(const PyLongObject *op)
+{
+    assert(PyLong_Check(op));
+    assert(PyUnstable_Long_IsCompact(op));
+    Py_ssize_t sign = 1 - (op->long_value.lv_tag & _PyLong_SIGN_MASK);
+    return sign * (Py_ssize_t)op->long_value.ob_digit[0];
+}
+
+#define PyUnstable_Long_CompactValue _PyLong_CompactValue
+
+
 #ifdef __cplusplus
 }
 #endif

Include/cpython/longobject.h

@@ -93,3 +93,8 @@ PyAPI_FUNC(PyObject *) _PyLong_GCD(PyObject *, PyObject *);
 
 PyAPI_FUNC(PyObject *) _PyLong_Rshift(PyObject *, size_t);
 PyAPI_FUNC(PyObject *) _PyLong_Lshift(PyObject *, size_t);
+
+
+PyAPI_FUNC(int) PyUnstable_Long_IsCompact(const PyLongObject* op);
+PyAPI_FUNC(Py_ssize_t) PyUnstable_Long_CompactValue(const PyLongObject* op);
+

Include/internal/pycore_long.h

@@ -118,6 +118,21 @@ PyAPI_FUNC(char*) _PyLong_FormatBytesWriter(
 #define SIGN_NEGATIVE 2
 #define NON_SIZE_BITS 3
 
+/* The functions _PyLong_IsCompact and _PyLong_CompactValue are defined
+ * in Include/cpython/longobject.h, since they need to be inline.
+ *
+ * "Compact" values have at least one bit to spare,
+ * so that addition and subtraction can be performed on the values
+ * without risk of overflow.
+ *
+ * The inline functions need tag bits.
+ * For readability, rather than do `#define SIGN_MASK _PyLong_SIGN_MASK`
+ * we define them to the numbers in both places and then assert that
+ * they're the same.
+ */
+static_assert(SIGN_MASK == _PyLong_SIGN_MASK, "SIGN_MASK does not match _PyLong_SIGN_MASK");
+static_assert(NON_SIZE_BITS == _PyLong_NON_SIZE_BITS, "NON_SIZE_BITS does not match _PyLong_NON_SIZE_BITS");
+
 /* All *compact" values are guaranteed to fit into
  * a Py_ssize_t with at least one bit to spare.
  * In other words, for 64 bit machines, compact
@@ -131,11 +146,6 @@ _PyLong_IsNonNegativeCompact(const PyLongObject* op) {
     return op->long_value.lv_tag <= (1 << NON_SIZE_BITS);
 }
 
-static inline int
-_PyLong_IsCompact(const PyLongObject* op) {
-    assert(PyLong_Check(op));
-    return op->long_value.lv_tag < (2 << NON_SIZE_BITS);
-}
 
 static inline int
 _PyLong_BothAreCompact(const PyLongObject* a, const PyLongObject* b) {
@@ -144,21 +154,6 @@ _PyLong_BothAreCompact(const PyLongObject* a, const PyLongObject* b) {
     return (a->long_value.lv_tag | b->long_value.lv_tag) < (2 << NON_SIZE_BITS);
 }
 
-/* Returns a *compact* value, iff `_PyLong_IsCompact` is true for `op`.
- *
- * "Compact" values have at least one bit to spare,
- * so that addition and subtraction can be performed on the values
- * without risk of overflow.
- */
-static inline Py_ssize_t
-_PyLong_CompactValue(const PyLongObject *op)
-{
-    assert(PyLong_Check(op));
-    assert(_PyLong_IsCompact(op));
-    Py_ssize_t sign = 1 - (op->long_value.lv_tag & SIGN_MASK);
-    return sign * (Py_ssize_t)op->long_value.ob_digit[0];
-}
-
 static inline bool
 _PyLong_IsZero(const PyLongObject *op)
 {

Lib/test/test_capi/test_long.py

@@ -0,0 +1,39 @@
+import unittest
+import sys
+
+from test.support import import_helper
+
+# Skip this test if the _testcapi module isn't available.
+_testcapi = import_helper.import_module('_testcapi')
+
+
+class LongTests(unittest.TestCase):
+
+    def test_compact(self):
+        for n in {
+            # Edge cases
+            *(2**n for n in range(66)),
+            *(-2**n for n in range(66)),
+            *(2**n - 1 for n in range(66)),
+            *(-2**n + 1 for n in range(66)),
+            # Essentially random
+            *(37**n for n in range(14)),
+            *(-37**n for n in range(14)),
+        }:
+            with self.subTest(n=n):
+                is_compact, value = _testcapi.call_long_compact_api(n)
+                if is_compact:
+                    self.assertEqual(n, value)
+
+    def test_compact_known(self):
+        # Sanity-check some implementation details (we don't guarantee
+        # that these are/aren't compact)
+        self.assertEqual(_testcapi.call_long_compact_api(-1), (True, -1))
+        self.assertEqual(_testcapi.call_long_compact_api(0), (True, 0))
+        self.assertEqual(_testcapi.call_long_compact_api(256), (True, 256))
+        self.assertEqual(_testcapi.call_long_compact_api(sys.maxsize),
+                         (False, -1))
+
+
+if __name__ == "__main__":
+    unittest.main()

Misc/NEWS.d/next/C API/2023-05-18-20-53-05.gh-issue-101291.ZBh9aR.rst

@@ -0,0 +1,3 @@
+Added unstable C API for extracting the value of "compact" integers:
+:c:func:`PyUnstable_Long_IsCompact` and
+:c:func:`PyUnstable_Long_CompactValue`.

Modules/_testcapi/long.c

@@ -534,6 +534,18 @@ test_long_numbits(PyObject *self, PyObject *Py_UNUSED(ignored))
     Py_RETURN_NONE;
 }
 
+static PyObject *
+check_long_compact_api(PyObject *self, PyObject *arg)
+{
+    assert(PyLong_Check(arg));
+    int is_compact = PyUnstable_Long_IsCompact((PyLongObject*)arg);
+    Py_ssize_t value = -1;
+    if (is_compact) {
+        value = PyUnstable_Long_CompactValue((PyLongObject*)arg);
+    }
+    return Py_BuildValue("in", is_compact, value);
+}
+
 static PyMethodDef test_methods[] = {
     {"test_long_and_overflow",  test_long_and_overflow,          METH_NOARGS},
     {"test_long_api",           test_long_api,                   METH_NOARGS},
@@ -543,6 +555,7 @@ static PyMethodDef test_methods[] = {
     {"test_long_long_and_overflow",test_long_long_and_overflow,  METH_NOARGS},
     {"test_long_numbits",       test_long_numbits,               METH_NOARGS},
     {"test_longlong_api",       test_longlong_api,               METH_NOARGS},
+    {"call_long_compact_api",   check_long_compact_api,          METH_O},
     {NULL},
 };
 

Objects/longobject.c

@@ -6366,3 +6366,17 @@ _PyLong_FiniTypes(PyInterpreterState *interp)
 {
     _PyStructSequence_FiniBuiltin(interp, &Int_InfoType);
 }
+
+#undef PyUnstable_Long_IsCompact
+
+int
+PyUnstable_Long_IsCompact(const PyLongObject* op) {
+    return _PyLong_IsCompact(op);
+}
+
+#undef PyUnstable_Long_CompactValue
+
+Py_ssize_t
+PyUnstable_Long_CompactValue(const PyLongObject* op) {
+    return _PyLong_CompactValue(op);
+}