added speedy BigInt extraction

newsfragments/3379.changed.md

@@ -0,0 +1 @@
+Sped up FromPyObject::extract for BigInt and BigUint by up to 43% (although mileage may vary depending on int size and sign)

src/conversions/num_bigint.rs

@@ -51,25 +51,37 @@ use crate::{
     ffi, types::*, FromPyObject, IntoPy, Py, PyAny, PyObject, PyResult, Python, ToPyObject,
 };
 
-use num_bigint::{BigInt, BigUint};
+use num_bigint::{BigInt, BigUint, Sign};
-use std::os::raw::{c_int, c_uchar};
+use std::os::raw::c_int;
 
 #[cfg(not(Py_LIMITED_API))]
-unsafe fn extract(ob: &PyLong, buffer: &mut [c_uchar], is_signed: c_int) -> PyResult<()> {
+use std::os::raw::c_uchar;
+
+#[cfg(Py_LIMITED_API)]
+use std::slice;
+
+#[cfg(not(Py_LIMITED_API))]
+#[inline]
+unsafe fn extract(ob: &PyLong, length: usize, is_signed: c_int) -> PyResult<Vec<u32>> {
+    let mut buffer = Vec::<u32>::with_capacity(length);
     crate::err::error_on_minusone(
         ob.py(),
         ffi::_PyLong_AsByteArray(
             ob.as_ptr() as *mut ffi::PyLongObject,
-            buffer.as_mut_ptr(),
+            buffer.as_mut_ptr() as *mut u8,
-            buffer.len(),
+            length * 4,
             1,
             is_signed,
         ),
-    )
+    )?;
+    buffer.set_len(length);
+
+    Ok(buffer)
 }
 
 #[cfg(Py_LIMITED_API)]
-unsafe fn extract(ob: &PyLong, buffer: &mut [c_uchar], is_signed: c_int) -> PyResult<()> {
+#[inline]
+unsafe fn extract(ob: &PyLong, length: usize, is_signed: c_int) -> PyResult<Vec<u32>> {
     use crate::intern;
     let py = ob.py();
     let kwargs = if is_signed != 0 {
@@ -81,14 +93,16 @@ unsafe fn extract(ob: &PyLong, buffer: &mut [c_uchar], is_signed: c_int) -> PyRe
     };
     let bytes_obj = ob
         .getattr(intern!(py, "to_bytes"))?
-        .call((buffer.len(), "little"), kwargs)?;
+        .call((length * 4, intern!(py, "little")), kwargs)?;
     let bytes: &PyBytes = bytes_obj.downcast_unchecked();
-    buffer.copy_from_slice(bytes.as_bytes());
+    let bytes_u32 = slice::from_raw_parts(bytes.as_bytes().as_ptr().cast(), length);
-    Ok(())
+
+    Ok(bytes_u32.to_vec())
 }
 
+// for identical functionality between BigInt and BigUint
 macro_rules! bigint_conversion {
-    ($rust_ty: ty, $is_signed: expr, $to_bytes: path, $from_bytes: path) => {
+    ($rust_ty: ty, $is_signed: expr, $to_bytes: path) => {
         #[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
         impl ToPyObject for $rust_ty {
             #[cfg(not(Py_LIMITED_API))]
@@ -129,62 +143,92 @@ macro_rules! bigint_conversion {
                 self.to_object(py)
             }
         }
+    };
+}
 
-        #[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
+bigint_conversion!(BigUint, 0, BigUint::to_bytes_le);
-        impl<'source> FromPyObject<'source> for $rust_ty {
+bigint_conversion!(BigInt, 1, BigInt::to_signed_bytes_le);
-            fn extract(ob: &'source PyAny) -> PyResult<$rust_ty> {
+
+#[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
+impl<'source> FromPyObject<'source> for BigInt {
+    fn extract(ob: &'source PyAny) -> PyResult<BigInt> {
         let py = ob.py();
         unsafe {
-                    let num: Py<PyLong> =
+            let num: Py<PyLong> = Py::from_owned_ptr_or_err(py, ffi::PyNumber_Index(ob.as_ptr()))?;
-                        Py::from_owned_ptr_or_err(py, ffi::PyNumber_Index(ob.as_ptr()))?;
+            let n_bits = {
-
-                    let n_bytes = {
                 cfg_if::cfg_if! {
                     if #[cfg(not(Py_LIMITED_API))] {
                         // fast path
                         let n_bits = ffi::_PyLong_NumBits(num.as_ptr());
                         if n_bits == (-1isize as usize) {
                             return Err(crate::PyErr::fetch(py));
-                                } else if n_bits == 0 {
-                                    0
-                                } else {
-                                    (n_bits - 1 + $is_signed) / 8 + 1
                         }
+                        n_bits
                     } else {
                         // slow path
                         let n_bits_obj = num.getattr(py, crate::intern!(py, "bit_length"))?.call0(py)?;
                         let n_bits_int: &PyLong = n_bits_obj.downcast_unchecked(py);
-                                let n_bits = n_bits_int.extract::<usize>()?;
+                        n_bits_int.extract::<usize>()?
-                                if n_bits == 0 {
-                                    0
-                                } else {
-                                    (n_bits - 1 + $is_signed) / 8 + 1
-                                }
                     }
                 }
             };
 
-                    if n_bytes <= 128 {
+            if n_bits == 0 {
-                        let mut buffer = [0; 128];
+                return Ok(BigInt::from(0isize));
-                        extract(num.as_ref(py), &mut buffer[..n_bytes], $is_signed)?;
+            }
-                        Ok($from_bytes(&buffer[..n_bytes]))
+            let n_digits = (n_bits + 32) / 32;
+            let mut buffer = extract(num.as_ref(py), n_digits, 1)?;
+            buffer
+                .iter_mut()
+                .for_each(|chunk| *chunk = u32::from_le(*chunk));
+
+            Ok(if buffer.last().unwrap() >> 31 != 0 {
+                buffer.iter_mut().for_each(|element| *element = !*element);
+                BigInt::new(Sign::Minus, buffer) - 1
             } else {
-                        let mut buffer = vec![0; n_bytes];
+                BigInt::new(Sign::Plus, buffer)
-                        extract(num.as_ref(py), &mut buffer, $is_signed)?;
+            })
-                        Ok($from_bytes(&buffer))
         }
     }
+}
+
+#[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
+impl<'source> FromPyObject<'source> for BigUint {
+    fn extract(ob: &'source PyAny) -> PyResult<BigUint> {
+        let py = ob.py();
+        unsafe {
+            let num: Py<PyLong> = Py::from_owned_ptr_or_err(py, ffi::PyNumber_Index(ob.as_ptr()))?;
+            let n_bits = {
+                cfg_if::cfg_if! {
+                    if #[cfg(not(Py_LIMITED_API))] {
+                        // fast path
+                        let n_bits = ffi::_PyLong_NumBits(num.as_ptr());
+                        if n_bits == (-1isize as usize) {
+                            return Err(crate::PyErr::fetch(py));
+                        }
+                        n_bits
+                    } else {
+                        // slow path
+                        let n_bits_obj = num.getattr(py, crate::intern!(py, "bit_length"))?.call0(py)?;
+                        let n_bits_int: &PyLong = n_bits_obj.downcast_unchecked(py);
+                        n_bits_int.extract::<usize>()?
                     }
                 }
             };
+
+            if n_bits == 0 {
+                return Ok(BigUint::from(0usize));
+            }
+            let n_digits = (n_bits + 31) / 32;
+            let mut buffer = extract(num.as_ref(py), n_digits, 0)?;
+            buffer
+                .iter_mut()
+                .for_each(|chunk| *chunk = u32::from_le(*chunk));
+
+            Ok(BigUint::new(buffer))
+        }
+    }
 }
-bigint_conversion!(BigUint, 0, BigUint::to_bytes_le, BigUint::from_bytes_le);
-bigint_conversion!(
-    BigInt,
-    1,
-    BigInt::to_signed_bytes_le,
-    BigInt::from_signed_bytes_le
-);
 
 #[cfg(test)]
 mod tests {
@@ -312,7 +356,7 @@ mod tests {
                 ($T:ty, $value:expr, $py:expr) => {
                     let value = $value;
                     println!("{}: {}", stringify!($T), value);
-                    let python_value = value.clone().to_object(py);
+                    let python_value = value.clone().into_py(py);
                     let roundtrip_value = python_value.extract::<$T>(py).unwrap();
                     assert_eq!(value, roundtrip_value);
                 };