--- 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 {}