rust/hg-cpython/src/pybytes_deref.rs
changeset 52163 7346f93be7a4
parent 49933 be3b545c5cff
child 52851 d9d6ae9b9722
--- a/rust/hg-cpython/src/pybytes_deref.rs	Wed Jun 19 17:03:13 2024 +0200
+++ b/rust/hg-cpython/src/pybytes_deref.rs	Wed Jun 19 19:10:49 2024 +0200
@@ -1,4 +1,7 @@
-use cpython::{PyBytes, Python};
+use crate::cpython::buffer::Element;
+use cpython::{
+    buffer::PyBuffer, exc::ValueError, PyBytes, PyErr, PyResult, Python,
+};
 use stable_deref_trait::StableDeref;
 
 /// Safe abstraction over a `PyBytes` together with the `&[u8]` slice
@@ -55,3 +58,67 @@
 // but here sending one to another thread is fine since we ensure it stays
 // valid.
 unsafe impl Send for PyBytesDeref {}
+
+///
+/// It also enables using a (wrapped) `PyBuffer` in GIL-unaware generic code.
+pub struct PyBufferDeref {
+    #[allow(unused)]
+    keep_alive: PyBuffer,
+
+    /// Borrows the buffer inside `self.keep_alive`,
+    /// but the borrow-checker cannot express self-referential structs.
+    data: *const [u8],
+}
+
+fn get_buffer<'a>(py: Python, buf: &'a PyBuffer) -> PyResult<&'a [u8]> {
+    let len = buf.item_count();
+
+    let cbuf = buf.buf_ptr();
+    let has_correct_item_size = std::mem::size_of::<u8>() == buf.item_size();
+    let is_valid_buffer = has_correct_item_size
+        && buf.is_c_contiguous()
+        && u8::is_compatible_format(buf.format())
+        && buf.readonly();
+
+    let bytes = if is_valid_buffer {
+        unsafe { std::slice::from_raw_parts(cbuf as *const u8, len) }
+    } else {
+        return Err(PyErr::new::<ValueError, _>(
+            py,
+            "Buffer has an invalid memory representation",
+        ));
+    };
+    Ok(bytes)
+}
+
+impl PyBufferDeref {
+    pub fn new(py: Python, buf: PyBuffer) -> PyResult<Self> {
+        Ok(Self {
+            data: get_buffer(py, &buf)?,
+            keep_alive: buf,
+        })
+    }
+}
+
+impl std::ops::Deref for PyBufferDeref {
+    type Target = [u8];
+
+    fn deref(&self) -> &[u8] {
+        // Safety: the raw pointer is valid as long as the PyBuffer is still
+        // alive, and the returned slice borrows `self`.
+        unsafe { &*self.data }
+    }
+}
+
+unsafe impl StableDeref for PyBufferDeref {}
+
+#[allow(unused)]
+fn static_assert_pybuffer_is_send() {
+    #[allow(clippy::no_effect)]
+    require_send::<PyBuffer>;
+}
+
+// Safety: PyBuffer is Send. Raw pointers are not by default,
+// but here sending one to another thread is fine since we ensure it stays
+// valid.
+unsafe impl Send for PyBufferDeref {}