diff --git a/newsfragments/6144.changed.md b/newsfragments/6144.changed.md new file mode 100644 index 00000000000..1df1e9e5efc --- /dev/null +++ b/newsfragments/6144.changed.md @@ -0,0 +1 @@ +Optimized `num-bigint` conversions on Python 3.14 and newer. diff --git a/src/conversions/num_bigint.rs b/src/conversions/num_bigint.rs index a5172e6ef3a..4ecc3d91e3d 100644 --- a/src/conversions/num_bigint.rs +++ b/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 { + digits: I, + acc: u64, + acc_bits: usize, + remaining: usize, +} + +#[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))] +impl Iterator for PyLongDigitIter +where + I: ExactSizeIterator, +{ + type Item = u32; + + fn next(&mut self) -> Option { + 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) { + (self.remaining, Some(self.remaining)) + } +} + +#[cfg(any(all(Py_3_14, not(Py_LIMITED_API)), Py_3_15))] +impl ExactSizeIterator for PyLongDigitIter +where + I: ExactSizeIterator, +{ + 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( + py: Python<'_>, + negative: bool, + bits: usize, + digits: I, +) -> Bound<'_, PyInt> +where + I: ExactSizeIterator, +{ + 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 { 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::(&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::(&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 { + 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(long: &Bound<'_, PyInt>) -> PyResult> { @@ -271,9 +459,7 @@ fn int_to_u32_vec(long: &Bound<'_, PyInt>) -> PyResult::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::().unwrap_err(); + assert!(err.is_instance_of::(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::().unwrap_err(); assert!(err.is_instance_of::(py)); diff --git a/src/conversions/std/num.rs b/src/conversions/std/num.rs index e44934bad33..13de965c2b5 100644 --- a/src/conversions/std/num.rs +++ b/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>( - py: Python<'py>, +pub(crate) fn pylong_from_digits>( + 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( 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::(); - 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::(); + + 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))]