Faster conversion from PyLong to u128/i128 with LIMITED_API

PyO3 · Dec 13, 2020 · 1a1ba74 · 1a1ba74
1 parent a93d97b
commit 1a1ba74
Showing 1 changed file with 30 additions and 34 deletions.
diff --git a/src/types/num.rs b/src/types/num.rs
@@ -216,24 +216,22 @@ mod fast_128bit_int_conversion {
     int_convert_128!(u128, 0);
 }
 
-// For ABI3 and PyPy, we implement the conversion (almost) manually.
+// 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};
     const BYTE_SIZE: usize = std::mem::size_of::<i128>();
+    const SHIFT: usize = 64;
 
-    fn split_128int(bytes: [u8; BYTE_SIZE]) -> ([u8; BYTE_SIZE / 2], [u8; BYTE_SIZE / 2]) {
-        let mut first = [0u8; BYTE_SIZE / 2];
-        let mut last = [0u8; BYTE_SIZE / 2];
-        first.copy_from_slice(&bytes[..BYTE_SIZE / 2]);
-        last.copy_from_slice(&bytes[BYTE_SIZE / 2..]);
-        (first, last)
+    fn slice_to_bytearr(bytes: &[u8]) -> [u8; BYTE_SIZE / 2] {
+        let mut res = [0u8; BYTE_SIZE / 2];
+        res.copy_from_slice(bytes);
+        res
     }
 
     // for 128bit Integers
     macro_rules! int_convert_128 {
-        ($rust_type: ty, $half_type: ty, $is_signed: literal) => {
+        ($rust_type: ty, $half_type: ty, $half_conversion: ident) => {
             impl ToPyObject for $rust_type {
                 #[inline]
                 fn to_object(&self, py: Python) -> PyObject {
@@ -243,19 +241,21 @@ mod slow_128bit_int_conversion {
 
             impl IntoPy<PyObject> for $rust_type {
                 fn into_py(self, py: Python) -> PyObject {
-                    let (first, last) = split_128int(self.to_le_bytes());
-                    let (first, last) =
-                        (u64::from_le_bytes(first), <$half_type>::from_le_bytes(last));
-                    let (first, last) = (first.into_py(py), last.into_py(py));
-                    let shift = 64u64.into_py(py);
+                    let le_bytes = self.to_le_bytes();
+                    let lower = u64::from_le_bytes(slice_to_bytearr(&le_bytes[..BYTE_SIZE / 2]));
+                    let upper =
+                        <$half_type>::from_le_bytes(slice_to_bytearr(&le_bytes[BYTE_SIZE / 2..]));
                     unsafe {
                         let shifted = PyObject::from_owned_ptr(
                             py,
-                            ffi::PyNumber_Lshift(last.as_ptr(), shift.as_ptr()),
+                            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(), first.as_ptr()),
+                            ffi::PyNumber_Or(shifted.as_ptr(), lower.into_py(py).as_ptr()),
                         )
                     }
                 }
@@ -264,30 +264,26 @@ mod slow_128bit_int_conversion {
             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; BYTE_SIZE];
-                    let kwargs = if $is_signed {
-                        Some([("signed", true)].into_py_dict(py))
-                    } else {
-                        None
-                    };
-                    let pybytes: &PyBytes = num
-                        .call_method("to_bytes", (bytes.len(), "little"), kwargs)?
-                        .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, i64, true);
-    int_convert_128!(u128, u64, false);
+    int_convert_128!(i128, i64, PyLong_AsLongLong);
+    int_convert_128!(u128, u64, PyLong_AsUnsignedLongLong);
 }
 
 #[cfg(test)]