luxolus
/
pyo3
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #1315 from PyO3/abi3-128bit-integer 

Faster conversion from u128/i128 to PyLong with abi3
This commit is contained in:
Yuji Kanagawa 2020-12-19 15:56:08 +09:00 committed by GitHub
parent 58ade65b99 cd7348fc51
commit e64dc12de1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 76 additions and 47 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.toml
View File

@ -32,6 +32,7 @@ hashbrown = { version = "0.9", optional = true }
assert_approx_eq = "1.1.0"
trybuild = "1.0.23"
rustversion = "1.0"
proptest = { version = "0.10.1", default-features = false, features = ["std"] }
[features]
default = ["macros"]

8
src/lib.rs
View File

@ -317,10 +317,10 @@ macro_rules! wrap_pymodule {
#[cfg(feature = "macros")]
macro_rules! py_run {
($py:expr, $($val:ident)+, $code:literal) => {{
pyo3::py_run_impl!($py, $($val)+, pyo3::indoc::indoc!($code))
$crate::py_run_impl!($py, $($val)+, $crate::indoc::indoc!($code))
}};
($py:expr, $($val:ident)+, $code:expr) => {{
pyo3::py_run_impl!($py, $($val)+, &pyo3::unindent::unindent($code))
$crate::py_run_impl!($py, $($val)+, &$crate::unindent::unindent($code))
}};
}
@ -329,8 +329,8 @@ macro_rules! py_run {
#[cfg(feature = "macros")]
macro_rules! py_run_impl {
($py:expr, $($val:ident)+, $code:expr) => {{
use pyo3::types::IntoPyDict;
use pyo3::ToPyObject;
use $crate::types::IntoPyDict;
use $crate::ToPyObject;
let d = [$((stringify!($val), $val.to_object($py)),)+].into_py_dict($py);
$py.run($code, None, Some(d))

114
src/types/num.rs
View File

@ -216,24 +216,15 @@ mod fast_128bit_int_conversion {
int_convert_128!(u128, 0);
}
// For ABI3 and PyPy, we implement the conversion using `call_method`.
// For ABI3 and PyPy, we implement the conversion manually.
#[cfg(any(Py_LIMITED_API, PyPy))]
mod slow_128bit_int_conversion {
use super::*;
use crate::types::{IntoPyDict, PyBytes, PyDict, PyType};
// Returns `{'signed': True}` if is_signed else None
fn kwargs(py: Python, is_signed: bool) -> Option<&PyDict> {
if is_signed {
Some([("signed", is_signed)].into_py_dict(py))
} else {
None
}
}
const SHIFT: usize = 64;
// for 128bit Integers
macro_rules! int_convert_128 {
($rust_type: ty, $is_signed: literal) => {
($rust_type: ty, $half_type: ty) => {
impl ToPyObject for $rust_type {
#[inline]
fn to_object(&self, py: Python) -> PyObject {
@ -243,43 +234,79 @@ mod slow_128bit_int_conversion {
impl IntoPy<PyObject> for $rust_type {
fn into_py(self, py: Python) -> PyObject {
let bytes = PyBytes::new(py, &self.to_le_bytes());
let longtype = unsafe { PyType::from_type_ptr(py, &mut ffi::PyLong_Type) };
longtype
.call_method("from_bytes", (bytes, "little"), kwargs(py, $is_signed))
.expect("Integer conversion (u128/i128 to PyLong) failed")
.into_py(py)
let lower = self as u64;
let upper = (self >> SHIFT) as $half_type;
unsafe {
let shifted = PyObject::from_owned_ptr(
py,
ffi::PyNumber_Lshift(
upper.into_py(py).as_ptr(),
SHIFT.into_py(py).as_ptr(),
),
);
PyObject::from_owned_ptr(
py,
ffi::PyNumber_Or(shifted.as_ptr(), lower.into_py(py).as_ptr()),
)
}
}
}
impl<'source> FromPyObject<'source> for $rust_type {
fn extract(ob: &'source PyAny) -> PyResult<$rust_type> {
let py = ob.py();
let num: &PyLong = unsafe {
crate::FromPyPointer::from_owned_ptr_or_err(
unsafe {
let lower = err_if_invalid_value(
py,
ffi::PyNumber_Index(ob.as_ptr()),
)
}?;
let mut bytes = [0u8; std::mem::size_of::<$rust_type>()];
let pybytes: &PyBytes = num
.call_method("to_bytes", (bytes.len(), "little"), kwargs(py, $is_signed))?
.extract()?;
bytes.copy_from_slice(pybytes.as_bytes());
Ok(<$rust_type>::from_le_bytes(bytes))
-1 as _,
ffi::PyLong_AsUnsignedLongLongMask(ob.as_ptr()),
)? as $rust_type;
let shifted = PyObject::from_owned_ptr(
py,
ffi::PyNumber_Rshift(ob.as_ptr(), SHIFT.into_py(py).as_ptr()),
);
let upper: $half_type = shifted.extract(py)?;
Ok((<$rust_type>::from(upper) << SHIFT) | lower)
}
}
}
};
}
int_convert_128!(i128, true);
int_convert_128!(u128, false);
int_convert_128!(i128, i64);
int_convert_128!(u128, u64);
}
#[cfg(test)]
mod test_128bit_intergers {
use super::*;
use proptest::prelude::*;
proptest! {
#[test]
fn test_i128_roundtrip(x: i128) {
Python::with_gil(|py| {
let x_py = x.into_py(py);
crate::py_run!(py, x_py, &format!("assert x_py == {}", x));
let roundtripped: i128 = x_py.extract(py).unwrap();
assert_eq!(x, roundtripped);
})
}
}
proptest! {
#[test]
fn test_u128_roundtrip(x: u128) {
Python::with_gil(|py| {
let x_py = x.into_py(py);
crate::py_run!(py, x_py, &format!("assert x_py == {}", x));
let roundtripped: u128 = x_py.extract(py).unwrap();
assert_eq!(x, roundtripped);
})
}
}
#[test]
fn test_i128_max() {
Python::with_gil(|py| {
@ -313,19 +340,20 @@ mod test_128bit_intergers {
}
#[test]
#[cfg(not(any(Py_LIMITED_API, PyPy)))]
fn test_u128_overflow() {
use crate::{exceptions, ffi, PyObject};
use std::os::raw::c_uchar;
fn test_i128_overflow() {
Python::with_gil(|py| {
let overflow_bytes: [c_uchar; 20] = [255; 20];
unsafe {
let obj =
ffi::_PyLong_FromByteArray(overflow_bytes.as_ptr() as *const c_uchar, 20, 1, 0);
let obj = PyObject::from_owned_ptr(py, obj);
let err = obj.extract::<u128>(py).unwrap_err();
assert!(err.is_instance::<exceptions::PyOverflowError>(py));
}
let obj = py.eval("(1 << 130) * -1", None, None).unwrap();
let err = obj.extract::<i128>().unwrap_err();
assert!(err.is_instance::<crate::exceptions::PyOverflowError>(py));
})
}
#[test]
fn test_u128_overflow() {
Python::with_gil(|py| {
let obj = py.eval("1 << 130", None, None).unwrap();
let err = obj.extract::<u128>().unwrap_err();
assert!(err.is_instance::<crate::exceptions::PyOverflowError>(py));
})
}
}