Skip to content

Commit

Permalink
added speedy BigInt extraction
Browse files Browse the repository at this point in the history
  • Loading branch information
@iliya-malecki @davidhewitt
iliya-malecki authored and davidhewitt committed Aug 17, 2023
1 parent 438d0a2 commit d1f0561
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 58 deletions.
1 change: 1 addition & 0 deletions newsfragments/3379.changed.md
View file
Original file line number Diff line number Diff line change
@@ -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)
160 changes: 102 additions & 58 deletions src/conversions/num_bigint.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,25 +51,37 @@ use crate::{
ffi, types::*, FromPyObject, IntoPy, Py, PyAny, PyObject, PyResult, Python, ToPyObject,
};

use num_bigint::{BigInt, BigUint};
use std::os::raw::{c_int, c_uchar};
use num_bigint::{BigInt, BigUint, Sign};
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.len(),
buffer.as_mut_ptr() as *mut u8,
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 {
Expand All @@ -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());
Ok(())
let bytes_u32 = slice::from_raw_parts(bytes.as_bytes().as_ptr().cast(), length);

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))]
Expand Down Expand Up @@ -129,62 +143,92 @@ macro_rules! bigint_conversion {
self.to_object(py)
}
}
};
}

#[cfg_attr(docsrs, doc(cfg(feature = "num-bigint")))]
impl<'source> FromPyObject<'source> for $rust_ty {
fn extract(ob: &'source PyAny) -> PyResult<$rust_ty> {
let py = ob.py();
unsafe {
let num: Py<PyLong> =
Py::from_owned_ptr_or_err(py, ffi::PyNumber_Index(ob.as_ptr()))?;
bigint_conversion!(BigUint, 0, BigUint::to_bytes_le);
bigint_conversion!(BigInt, 1, BigInt::to_signed_bytes_le);

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
}
} 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>()?;
if n_bits == 0 {
0
} else {
(n_bits - 1 + $is_signed) / 8 + 1
}
}
#[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> = 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(BigInt::from(0isize));
}
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 {
BigInt::new(Sign::Plus, buffer)
})
}
}
}

if n_bytes <= 128 {
let mut buffer = [0; 128];
extract(num.as_ref(py), &mut buffer[..n_bytes], $is_signed)?;
Ok($from_bytes(&buffer[..n_bytes]))
#[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 {
let mut buffer = vec![0; n_bytes];
extract(num.as_ref(py), &mut buffer, $is_signed)?;
Ok($from_bytes(&buffer))
// 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 {
Expand Down Expand Up @@ -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);
};
Expand Down

0 comments on commit d1f0561

Please sign in to comment.