bytes specialization via unsafe slice cast

PyO3 · Aug 7, 2024 · 6ca4fe4 · 6ca4fe4
1 parent e532e0d
commit 6ca4fe4
Show file tree

Hide file tree

Showing 7 changed files with 85 additions and 35 deletions.
diff --git a/src/conversion.rs b/src/conversion.rs
@@ -4,7 +4,7 @@ use crate::err::PyResult;
 use crate::inspect::types::TypeInfo;
 use crate::pyclass::boolean_struct::False;
 use crate::types::any::PyAnyMethods;
-use crate::types::PyTuple;
+use crate::types::{PyBytes, PyTuple};
 use crate::{
     ffi, Borrowed, Bound, BoundObject, Py, PyAny, PyClass, PyErr, PyObject, PyRef, PyRefMut, Python,
 };
@@ -202,15 +202,16 @@ pub trait IntoPyObject<'py>: Sized {
     fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error>;
 
     /// Converts slice of self into a Python object
-    fn iter_into_pyobject<I, E>(
+    fn iter_into_pyobject<I>(
         iter: I,
         py: Python<'py>,
-        _: private::Token,
+        _: crate::conversion::private::Token,
     ) -> Result<Bound<'py, PyAny>, PyErr>
     where
-        I: IntoIterator<Item = Self, IntoIter = E>,
-        E: ExactSizeIterator<Item = Self>,
-        PyErr: From<Self::Error>,
+        I: crate::conversion::SliceableIntoPyObjectIterator<Item = Self>,
+        I::Item: IntoPyObject<'py>,
+        I::IntoIter: ExactSizeIterator,
+        PyErr: From<<I::Item as IntoPyObject<'py>>::Error>,
     {
         let mut iter = iter.into_iter().map(|e| {
             e.into_pyobject(py)
@@ -223,6 +224,10 @@ pub trait IntoPyObject<'py>: Sized {
     }
 }
 
+pub trait SliceableIntoPyObjectIterator: IntoIterator {
+    fn as_bytes_slice(&self) -> &[u8];
+}
+
 pub(crate) mod private {
     pub struct Token;
 }

diff --git a/src/conversions/smallvec.rs b/src/conversions/smallvec.rs
@@ -70,11 +70,22 @@ where
     ///
     /// [`PyBytes`]: crate::types::PyBytes
     /// [`PyList`]: crate::types::PyList
+    #[inline]
     fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
         <A::Item>::iter_into_pyobject(self, py, crate::conversion::private::Token)
     }
 }
 
+impl<A> crate::conversion::SliceableIntoPyObjectIterator for SmallVec<A>
+where
+    A: Array,
+{
+    #[inline]
+    fn as_bytes_slice(&self) -> &[u8] {
+        unsafe { std::slice::from_raw_parts(self.as_ptr().cast::<u8>(), self.len()) }
+    }
+}
+
 impl<'py, A> FromPyObject<'py> for SmallVec<A>
 where
     A: Array,

diff --git a/src/conversions/std/array.rs b/src/conversions/std/array.rs
@@ -55,6 +55,12 @@ where
     }
 }
 
+impl<T, const N: usize> crate::conversion::SliceableIntoPyObjectIterator for [T; N] {
+    fn as_bytes_slice(&self) -> &[u8] {
+        unsafe { std::slice::from_raw_parts(self.as_ptr().cast::<u8>(), self.len()) }
+    }
+}
+
 impl<T, const N: usize> ToPyObject for [T; N]
 where
     T: ToPyObject,

diff --git a/src/conversions/std/num.rs b/src/conversions/std/num.rs
@@ -213,31 +213,21 @@ impl<'py> IntoPyObject<'py> for u8 {
         }
     }
 
-    fn iter_into_pyobject<I, E>(
+    #[inline]
+    fn iter_into_pyobject<I>(
         iter: I,
         py: Python<'py>,
         _: crate::conversion::private::Token,
     ) -> Result<Bound<'py, PyAny>, PyErr>
     where
-        I: IntoIterator<Item = Self, IntoIter = E>,
-        E: ExactSizeIterator<Item = Self>,
+        I: crate::conversion::SliceableIntoPyObjectIterator<Item = Self>,
+        I::Item: IntoPyObject<'py>,
+        I::IntoIter: ExactSizeIterator,
     {
-        let mut iter = iter.into_iter();
-        let len = iter.len();
-
-        PyBytes::new_with(py, len, |buf| {
-            let mut counter = 0;
-            for (slot, byte) in buf.iter_mut().zip(&mut iter) {
-                *slot = byte;
-                counter += 1;
-            }
-
-            assert!(iter.next().is_none(), "Attempted to create PyBytes but `iter` was larger than reported by its `ExactSizeIterator` implementation.");
-            assert_eq!(len, counter, "Attempted to create PyBytes but `iter` was smaller than reported by its `ExactSizeIterator` implementation.");
-
-            Ok(())
-        })
-        .map(Bound::into_any)
+        // we can make use of concrete code relating to u8 in here, BUT
+        // we have a problem due to &u8 vs u8, which makes it difficult to get
+        // a slice out without some crazy specialization
+        Ok(PyBytes::new(py, iter.as_bytes_slice()).into_any())
     }
 }
 
@@ -250,21 +240,21 @@ impl<'py> IntoPyObject<'py> for &u8 {
         (*self).into_pyobject(py)
     }
 
-    fn iter_into_pyobject<I, E>(
+    #[inline]
+    fn iter_into_pyobject<I>(
         iter: I,
         py: Python<'py>,
         _: crate::conversion::private::Token,
     ) -> Result<Bound<'py, PyAny>, PyErr>
     where
-        I: IntoIterator<Item = Self, IntoIter = E>,
-        E: ExactSizeIterator<Item = Self>,
-        PyErr: From<Self::Error>,
+        I: crate::conversion::SliceableIntoPyObjectIterator<Item = Self>,
+        I::Item: IntoPyObject<'py>,
+        I::IntoIter: ExactSizeIterator,
     {
-        u8::iter_into_pyobject(
-            iter.into_iter().copied(),
-            py,
-            crate::conversion::private::Token,
-        )
+        // we can make use of concrete code relating to u8 in here, BUT
+        // we have a problem due to &u8 vs u8, which makes it difficult to get
+        // a slice out without some crazy specialization
+        Ok(PyBytes::new(py, iter.as_bytes_slice()).into_any())
     }
 }
 

diff --git a/src/conversions/std/slice.rs b/src/conversions/std/slice.rs
@@ -32,11 +32,19 @@ where
     ///
     /// [`PyBytes`]: crate::types::PyBytes
     /// [`PyList`]: crate::types::PyList
+    #[inline]
     fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
         <&T>::iter_into_pyobject(self, py, crate::conversion::private::Token)
     }
 }
 
+impl<'a, T> crate::conversion::SliceableIntoPyObjectIterator for &'a [T] {
+    #[inline]
+    fn as_bytes_slice(&self) -> &[u8] {
+        unsafe { std::slice::from_raw_parts(self.as_ptr().cast::<u8>(), self.len()) }
+    }
+}
+
 impl<'a> crate::conversion::FromPyObjectBound<'a, '_> for &'a [u8] {
     fn from_py_object_bound(obj: crate::Borrowed<'a, '_, PyAny>) -> PyResult<Self> {
         Ok(obj.downcast::<PyBytes>()?.as_bytes())
@@ -91,8 +99,9 @@ where
     type Output = Bound<'py, Self::Target>;
     type Error = PyErr;
 
+    #[inline]
     fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
-        <&T>::iter_into_pyobject(self.iter(), py, crate::conversion::private::Token)
+        <&T>::iter_into_pyobject(&*self, py, crate::conversion::private::Token)
     }
 }
 

diff --git a/src/conversions/std/vec.rs b/src/conversions/std/vec.rs
@@ -53,11 +53,19 @@ where
     ///
     /// [`PyBytes`]: crate::types::PyBytes
     /// [`PyList`]: crate::types::PyList
+    #[inline]
     fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
         T::iter_into_pyobject(self, py, crate::conversion::private::Token)
     }
 }
 
+impl<T> crate::conversion::SliceableIntoPyObjectIterator for Vec<T> {
+    #[inline]
+    fn as_bytes_slice(&self) -> &[u8] {
+        unsafe { std::slice::from_raw_parts(self.as_ptr().cast::<u8>(), self.len()) }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::conversion::IntoPyObject;

diff --git a/src/types/bytes.rs b/src/types/bytes.rs
@@ -2,6 +2,7 @@ use crate::ffi_ptr_ext::FfiPtrExt;
 use crate::instance::{Borrowed, Bound};
 use crate::types::any::PyAnyMethods;
 use crate::{ffi, Py, PyAny, PyResult, Python};
+use std::mem::MaybeUninit;
 use std::ops::Index;
 use std::slice::SliceIndex;
 use std::str;
@@ -113,6 +114,26 @@ impl PyBytes {
         }
     }
 
+    pub(crate) fn new_with_uninit<F>(
+        py: Python<'_>,
+        len: usize,
+        init: F,
+    ) -> PyResult<Bound<'_, PyBytes>>
+    where
+        F: FnOnce(&mut [MaybeUninit<u8>]) -> PyResult<()>,
+    {
+        unsafe {
+            let pyptr = ffi::PyBytes_FromStringAndSize(std::ptr::null(), len as ffi::Py_ssize_t);
+            // Check for an allocation error and return it
+            let pybytes = pyptr.assume_owned_or_err(py)?.downcast_into_unchecked();
+            let buffer: *mut MaybeUninit<u8> = ffi::PyBytes_AsString(pyptr).cast();
+            debug_assert!(!buffer.is_null());
+            // (Further) Initialise the bytestring in init
+            // If init returns an Err, pypybytearray will automatically deallocate the buffer
+            init(std::slice::from_raw_parts_mut(buffer, len)).map(|_| pybytes)
+        }
+    }
+
     /// Deprecated name for [`PyBytes::new_with`].
     #[deprecated(since = "0.23.0", note = "renamed to `PyBytes::new_with`")]
     #[inline]