perf: use new PyLong* API for num-bigint feature

newsfragments/6144.changed.md

@@ -0,0 +1 @@
+Optimized `num-bigint` conversions on Python 3.14 and newer.

src/conversions/num_bigint.rs

@@ -49,6 +49,10 @@
 //! assert n + 1 == value
 //! ```
 
+#[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+use crate::conversions::std::num::{
+    is_30bit_layout, pylong_from_digits, pylong_visit_digits, PYLONG_BITS_IN_DIGIT,
+};
 #[allow(unused_imports, reason = "conditionally used")]
 use crate::platform::prelude::*;
 #[cfg(Py_LIMITED_API)]
@@ -60,16 +64,91 @@ use crate::{
 
 use num_bigint::{BigInt, BigUint};
 
+#[cfg(any(not(Py_LIMITED_API), Py_3_15))]
+use num_bigint::Sign;
+
 #[cfg(feature = "experimental-inspect")]
 use crate::inspect::PyStaticExpr;
 #[cfg(feature = "experimental-inspect")]
 use crate::PyTypeInfo;
-#[cfg(not(Py_LIMITED_API))]
-use num_bigint::Sign;
+
+#[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+struct PyLongDigitIter<I> {
+    digits: I,
+    acc: u64,
+    acc_bits: usize,
+    remaining: usize,
+}
+
+#[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+impl<I> Iterator for PyLongDigitIter<I>
+where
+    I: ExactSizeIterator<Item = u32>,
+{
+    type Item = u32;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        const MASK: u32 = (1 << PYLONG_BITS_IN_DIGIT) - 1;
+
+        if self.remaining == 0 {
+            return None;
+        }
+
+        while self.acc_bits < PYLONG_BITS_IN_DIGIT {
+            let Some(digit) = self.digits.next() else {
+                break;
+            };
+
+            self.acc |= u64::from(digit) << self.acc_bits;
+            self.acc_bits += u32::BITS as usize;
+        }
+
+        let digit = self.acc as u32 & MASK;
+        self.acc >>= PYLONG_BITS_IN_DIGIT;
+        self.acc_bits = self.acc_bits.saturating_sub(PYLONG_BITS_IN_DIGIT);
+        self.remaining -= 1;
+        Some(digit)
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.remaining, Some(self.remaining))
+    }
+}
+
+#[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+impl<I> ExactSizeIterator for PyLongDigitIter<I>
+where
+    I: ExactSizeIterator<Item = u32>,
+{
+    fn len(&self) -> usize {
+        self.remaining
+    }
+}
+
+#[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+#[inline]
+fn pylong_from_u32_digits<I>(
+    py: Python<'_>,
+    negative: bool,
+    bits: usize,
+    digits: I,
+) -> Bound<'_, PyInt>
+where
+    I: ExactSizeIterator<Item = u32>,
+{
+    let py_digits_len = bits.div_ceil(PYLONG_BITS_IN_DIGIT).max(1);
+    let digits = PyLongDigitIter {
+        digits,
+        acc: 0,
+        acc_bits: 0,
+        remaining: py_digits_len,
+    };
+    pylong_from_digits(py, negative, digits)
+}
 
 // for identical functionality between BigInt and BigUint
 macro_rules! bigint_conversion {
-    ($rust_ty: ty, $is_signed: literal) => {
+    ($rust_ty: ty, $is_signed: literal, $bits:path, $iter_u32_digits:path, $negative:expr) => {
         #[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
         impl<'py> IntoPyObject<'py> for $rust_ty {
             type Target = PyInt;
@@ -97,6 +176,22 @@ macro_rules! bigint_conversion {
             fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
                 use num_traits::ToBytes;
 
+                #[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+                {
+                    if is_30bit_layout() {
+                        let bits: usize = $bits(self)
+                            .try_into()
+                            .expect(concat!(stringify!($rust_ty), " bit length fits in usize"));
+
+                        return Ok(pylong_from_u32_digits(
+                            py,
+                            $negative(self),
+                            bits,
+                            $iter_u32_digits(self),
+                        ));
+                    }
+                }
+
                 #[cfg(all(not(Py_LIMITED_API), Py_3_13))]
                 {
                     use crate::conversions::std::num::int_from_ne_bytes;
@@ -135,8 +230,20 @@ macro_rules! bigint_conversion {
     };
 }
 
-bigint_conversion!(BigUint, false);
-bigint_conversion!(BigInt, true);
+bigint_conversion!(
+    BigUint,
+    false,
+    BigUint::bits,
+    BigUint::iter_u32_digits,
+    |_| false
+);
+bigint_conversion!(
+    BigInt,
+    true,
+    BigInt::bits,
+    BigInt::iter_u32_digits,
+    |value: &BigInt| BigInt::sign(value) == Sign::Minus
+);
 
 #[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
 impl<'py> FromPyObject<'_, 'py> for BigInt {
@@ -154,6 +261,16 @@ impl<'py> FromPyObject<'_, 'py> for BigInt {
             num_owned = nb_index(&ob)?;
             num_owned.as_borrowed()
         };
+        #[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+        if is_30bit_layout() {
+            return pylong_visit_digits(num.as_any().as_borrowed(), |negative, compact, digits| {
+                let Some(digits) = digits else {
+                    return Ok(BigInt::from(compact));
+                };
+                let sign = if negative { Sign::Minus } else { Sign::Plus };
+                Ok(BigInt::new(sign, int_from_pylong_digits(digits)))
+            });
+        }
         #[cfg(not(Py_LIMITED_API))]
         {
             let mut buffer = int_to_u32_vec::<true>(&num)?;
@@ -206,6 +323,20 @@ impl<'py> FromPyObject<'_, 'py> for BigUint {
             num_owned = nb_index(&ob)?;
             num_owned.as_borrowed()
         };
+        #[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+        if is_30bit_layout() {
+            return pylong_visit_digits(num.as_any().as_borrowed(), |negative, compact, digits| {
+                if negative {
+                    return Err(crate::exceptions::PyValueError::new_err(
+                        "can't convert negative int to unsigned",
+                    ));
+                }
+                let Some(digits) = digits else {
+                    return Ok(BigUint::from(compact as u64));
+                };
+                Ok(BigUint::new(int_from_pylong_digits(digits)))
+            });
+        }
         #[cfg(not(Py_LIMITED_API))]
         {
             let buffer = int_to_u32_vec::<false>(&num)?;
@@ -223,6 +354,63 @@ impl<'py> FromPyObject<'_, 'py> for BigUint {
     }
 }
 
+#[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+#[inline]
+fn int_from_pylong_digits(digits: &[u32]) -> Vec<u32> {
+    let n_digits = digits
+        .last()
+        .map(|last| {
+            let total_bits = (digits.len() - 1) * PYLONG_BITS_IN_DIGIT
+                + (u32::BITS - last.leading_zeros()) as usize;
+
+            total_bits.div_ceil(u32::BITS as usize)
+        })
+        .unwrap_or(0);
+
+    let mut py_digits = Vec::with_capacity(n_digits);
+    let ptr: *mut u32 = py_digits.as_mut_ptr();
+
+    if let Some((&last, init)) = digits.split_last() {
+        let mut acc = 0;
+        let mut acc_bits = 0;
+        let mut written = 0;
+
+        for &digit in init {
+            acc |= u64::from(digit) << acc_bits;
+            let new_bits = acc_bits + PYLONG_BITS_IN_DIGIT;
+
+            if new_bits >= u32::BITS as usize {
+                unsafe {
+                    // SAFETY: the total number of writes is bounded by `n_digits`
+                    ptr.add(written).write(acc as u32);
+                    written += 1;
+                }
+                acc >>= u32::BITS;
+                acc_bits = new_bits - u32::BITS as usize;
+            } else {
+                acc_bits = new_bits;
+            }
+        }
+
+        acc |= u64::from(last) << acc_bits;
+        while written < n_digits {
+            unsafe {
+                // SAFETY: `written < n_digits <= capacity` by construction
+                ptr.add(written).write(acc as u32);
+                written += 1;
+            }
+            acc >>= u32::BITS;
+        }
+
+        unsafe {
+            // SAFETY: exactly `written` elements were initialized above
+            py_digits.set_len(written);
+        }
+    }
+
+    py_digits
+}
+
 #[cfg(not(any(Py_LIMITED_API, Py_3_13)))]
 #[inline]
 fn int_to_u32_vec<const SIGNED: bool>(long: &Bound<'_, PyInt>) -> PyResult<Vec<u32>> {
@@ -271,9 +459,7 @@ fn int_to_u32_vec<const SIGNED: bool>(long: &Bound<'_, PyInt>) -> PyResult<Vec<u
     }
     let n_bytes =
         unsafe { ffi::PyLong_AsNativeBytes(long.as_ptr().cast(), core::ptr::null_mut(), 0, flags) };
-    let n_bytes_unsigned: usize = n_bytes
-        .try_into()
-        .map_err(|_| crate::PyErr::fetch(long.py()))?;
+    let n_bytes_unsigned: usize = n_bytes.try_into().map_err(|_| PyErr::fetch(long.py()))?;
     if n_bytes == 0 {
         return Ok(buffer);
     }
@@ -432,6 +618,41 @@ class C:
         })
     }
 
+    #[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+    #[test]
+    fn pylong_export() {
+        Python::attach(|py| {
+            if !is_30bit_layout() {
+                return;
+            }
+
+            let small = 42_i32.into_pyobject(py).unwrap();
+            let (negative, compact, digits) =
+                pylong_visit_digits(small.as_any().as_borrowed(), |negative, compact, digits| {
+                    Ok((negative, compact, digits.map(<[u32]>::to_vec)))
+                })
+                .unwrap();
+            assert!(!negative);
+            assert_eq!(compact, 42);
+            assert_eq!(digits, None);
+
+            let big = BigInt::new(Sign::Minus, vec![u32::MAX, 0x8000_0001, 0x1234_5678, 1])
+                .into_pyobject(py)
+                .unwrap();
+            let (negative, digits) =
+                pylong_visit_digits(big.as_any().as_borrowed(), |negative, _, digits| {
+                    Ok((negative, digits.map(<[u32]>::to_vec)))
+                })
+                .unwrap();
+            assert!(negative);
+            let digits = digits.unwrap();
+            assert_eq!(
+                int_from_pylong_digits(&digits),
+                vec![u32::MAX, 0x8000_0001, 0x1234_5678, 1]
+            );
+        });
+    }
+
     /// `OverflowError` on converting Python int to BigInt, see issue #629
     #[test]
     fn check_overflow() {
@@ -464,10 +685,24 @@ class C:
         });
     }
 
+    #[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
+    #[test]
+    fn biguint_negative() {
+        Python::attach(|py| {
+            let value = py.eval(c"-(1 << 130)", None, None).unwrap();
+            let err = value.extract::<BigUint>().unwrap_err();
+            assert!(err.is_instance_of::<crate::exceptions::PyValueError>(py));
+            assert_eq!(
+                err.value(py).to_string(),
+                "can't convert negative int to unsigned"
+            );
+        });
+    }
+
     #[test]
     fn from_py_float_type_error() {
         Python::attach(|py| {
-            let obj = (12.3f64).into_pyobject(py).unwrap();
+            let obj = 12.3f64.into_pyobject(py).unwrap();
             let err = obj.extract::<BigInt>().unwrap_err();
             assert!(err.is_instance_of::<PyTypeError>(py));
 

src/conversions/std/num.rs

@@ -394,11 +394,11 @@ impl Drop for ExportGuard {
 // Builds an int from an iterator of 30-bit digits
 #[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))]
 #[inline]
-pub(crate) fn pylong_from_digits<'py, I: ExactSizeIterator<Item = u32>>(
-    py: Python<'py>,
+pub(crate) fn pylong_from_digits<I: ExactSizeIterator<Item = u32>>(
+    py: Python<'_>,
     negative: bool,
     digits: I,
-) -> Bound<'py, PyInt> {
+) -> Bound<'_, PyInt> {
     let digits_len = digits.len();
     let mut ptr = core::ptr::null_mut();
     let writer = unsafe {
@@ -430,19 +430,24 @@ pub(crate) fn pylong_visit_digits<R>(
             ffi::PyLong_Export(obj.as_ptr(), long_export.as_mut_ptr()),
         )?;
     }
-    let export_guard = ExportGuard(unsafe { long_export.assume_init() });
-    let long_export_ref = &export_guard.0;
-    let value = long_export_ref.value;
-    if long_export_ref.digits.is_null() {
-        let negative = long_export_ref.value < 0;
-        f(negative, value, None)
-    } else {
-        let negative = long_export_ref.negative != 0;
-        let n_digits = long_export_ref.ndigits as usize;
-        let ptr = long_export_ref.digits.cast::<u32>();
-        let digits = unsafe { core::slice::from_raw_parts(ptr, n_digits) };
-        f(negative, value, Some(digits))
+    let long_export = unsafe { long_export.assume_init() };
+    let ptr = long_export.digits.cast::<u32>();
+
+    if ptr.is_null() {
+        // `value` is only valid when `digits` is NULL, and `PyLong_FreeExport()`
+        // is optional in that case
+        //
+        // See: https://docs.python.org/3/c-api/long.html#c.PyLong_FreeExport
+        return f(long_export.value < 0, long_export.value, None);
     }
+    // Keep the export alive while `digits` borrows the exported buffer
+    let export_guard = ExportGuard(long_export);
+
+    let negative = export_guard.0.negative != 0;
+    let n_digits = export_guard.0.ndigits as usize;
+    let digits = unsafe { core::slice::from_raw_parts(ptr, n_digits) };
+
+    f(negative, 0, Some(digits))
 }
 
 #[cfg(any(not(Py_LIMITED_API), Py_3_15))]