diff -r 13815c9decd4 -r 7346f93be7a4 rust/hg-cpython/src/pybytes_deref.rs --- 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::() == 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::( + py, + "Buffer has an invalid memory representation", + )); + }; + Ok(bytes) +} + +impl PyBufferDeref { + pub fn new(py: Python, buf: PyBuffer) -> PyResult { + 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::; +} + +// 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 {}