Skip to content

Commit fd029dc

Browse files
youknowoneyouknowone
committed
PyBuffer overhaul
PyBuffer doesn't require vtable access for options maintaining internal using manual Rc model
1 parent dcc8043 commit fd029dc

File tree

12 files changed

+300
-377
lines changed

12 files changed

+300
-377
lines changed

‎vm/src/buffer.rs

Lines changed: 53 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,10 @@
33
use crate::common::borrow::{BorrowedValue, BorrowedValueMut};
44
use crate::common::rc::PyRc;
55
use crate::PyThreadingConstraint;
6-
use crate::{PyObjectRef, PyResult, TypeProtocol};
7-
use crate::{TryFromBorrowedObject, VirtualMachine};
8-
use std::{borrow::Cow, fmt::Debug, ops::Deref};
6+
use crate::{PyObjectRef, PyResult, TryFromBorrowedObject, TypeProtocol, VirtualMachine};
7+
use std::{borrow::Cow, fmt::Debug};
98

10-
pub trait PyBuffer: Debug + PyThreadingConstraint {
11-
fn get_options(&self) -> &BufferOptions;
9+
pub trait PyBufferInternal: Debug + PyThreadingConstraint {
1210
/// Get the full inner buffer of this memory. You probably want [`as_contiguous()`], as
1311
/// `obj_bytes` doesn't take into account the range a memoryview might operate on, among other
1412
/// footguns.
@@ -18,49 +16,78 @@ pub trait PyBuffer: Debug + PyThreadingConstraint {
1816
/// might operate on, among other footguns.
1917
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]>;
2018
fn release(&self);
19+
// not included in PyBuffer protocol itself
20+
fn retain(&self);
21+
}
22+
23+
#[derive(Debug)]
24+
pub struct PyBuffer {
25+
pub options: BufferOptions,
26+
pub(crate) internal: PyRc<dyn PyBufferInternal>,
27+
}
2128

22-
fn as_contiguous(&self) -> Option<BorrowedValue<[u8]>> {
23-
if !self.get_options().contiguous {
29+
impl PyBuffer {
30+
pub fn new(buffer: impl PyBufferInternal + 'static, options: BufferOptions) -> Self {
31+
buffer.retain();
32+
Self {
33+
options,
34+
internal: PyRc::new(buffer),
35+
}
36+
}
37+
pub fn as_contiguous(&self) -> Option<BorrowedValue<[u8]>> {
38+
if !self.options.contiguous {
2439
return None;
2540
}
26-
Some(self.obj_bytes())
41+
Some(self.internal.obj_bytes())
2742
}
2843

29-
fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
30-
if !self.get_options().contiguous {
44+
pub fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
45+
if !self.options.contiguous {
3146
return None;
3247
}
33-
Some(self.obj_bytes_mut())
48+
Some(self.internal.obj_bytes_mut())
3449
}
3550

36-
fn to_contiguous(&self) -> Vec<u8> {
37-
self.obj_bytes().to_vec()
51+
pub fn to_contiguous(&self) -> Vec<u8> {
52+
self.internal.obj_bytes().to_vec()
53+
}
54+
55+
pub fn clone_with_options(&self, options: BufferOptions) -> Self {
56+
self.internal.retain();
57+
Self {
58+
options,
59+
internal: self.internal.clone(),
60+
}
3861
}
3962
}
4063

4164
#[derive(Debug, Clone)]
4265
pub struct BufferOptions {
43-
pub readonly: bool,
66+
// buf
67+
// obj
4468
pub len: usize,
69+
pub readonly: bool,
4570
pub itemsize: usize,
46-
pub contiguous: bool,
4771
pub format: Cow<'static, str>,
48-
// TODO: support multiple dimension array
49-
pub ndim: usize,
72+
pub ndim: usize, // TODO: support multiple dimension array
5073
pub shape: Vec<usize>,
5174
pub strides: Vec<isize>,
75+
// suboffsets
76+
77+
// RustPython fields
78+
pub contiguous: bool,
5279
}
5380

5481
impl BufferOptions {
5582
pub const DEFAULT: Self = BufferOptions {
56-
readonly: true,
5783
len: 0,
84+
readonly: true,
5885
itemsize: 1,
59-
contiguous: true,
6086
format: Cow::Borrowed("B"),
6187
ndim: 1,
6288
shape: Vec::new(),
6389
strides: Vec::new(),
90+
contiguous: true,
6491
};
6592
}
6693

@@ -70,15 +97,12 @@ impl Default for BufferOptions {
7097
}
7198
}
7299

73-
#[derive(Debug)]
74-
pub struct PyBufferRef(Box<dyn PyBuffer>);
75-
76-
impl TryFromBorrowedObject for PyBufferRef {
100+
impl TryFromBorrowedObject for PyBuffer {
77101
fn try_from_borrowed_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResult<Self> {
78102
let obj_cls = obj.class();
79103
for cls in obj_cls.iter_mro() {
80104
if let Some(f) = cls.slots.as_buffer.as_ref() {
81-
return f(obj, vm).map(|x| PyBufferRef(x));
105+
return f(obj, vm);
82106
}
83107
}
84108
Err(vm.new_type_error(format!(
@@ -88,76 +112,15 @@ impl TryFromBorrowedObject for PyBufferRef {
88112
}
89113
}
90114

91-
impl Drop for PyBufferRef {
92-
fn drop(&mut self) {
93-
self.0.release();
94-
}
95-
}
96-
97-
impl Deref for PyBufferRef {
98-
type Target = dyn PyBuffer;
99-
100-
fn deref(&self) -> &Self::Target {
101-
self.0.deref()
102-
}
103-
}
104-
105-
impl PyBufferRef {
106-
pub fn new(buffer: impl PyBuffer + 'static) -> Self {
107-
Self(Box::new(buffer))
108-
}
109-
110-
pub fn into_rcbuf(self) -> RcBuffer {
111-
// move self.0 out of self; PyBufferRef impls Drop so it's tricky
112-
let this = std::mem::ManuallyDrop::new(self);
113-
let buf_box = unsafe { std::ptr::read(&this.0) };
114-
RcBuffer(buf_box.into())
115-
}
116-
}
117-
118-
impl From<Box<dyn PyBuffer>> for PyBufferRef {
119-
fn from(buffer: Box<dyn PyBuffer>) -> Self {
120-
PyBufferRef(buffer)
121-
}
122-
}
123-
124-
#[derive(Debug, Clone)]
125-
pub struct RcBuffer(PyRc<dyn PyBuffer>);
126-
impl Deref for RcBuffer {
127-
type Target = dyn PyBuffer;
128-
fn deref(&self) -> &Self::Target {
129-
self.0.deref()
130-
}
131-
}
132-
133-
impl Drop for RcBuffer {
115+
impl Drop for PyBuffer {
134116
fn drop(&mut self) {
135-
// check if this is the last rc before the inner buffer gets dropped
136-
if let Some(buf) = PyRc::get_mut(&mut self.0) {
137-
buf.release()
138-
}
117+
self.internal.release();
139118
}
140119
}
141120

142-
impl PyBuffer for RcBuffer {
143-
fn get_options(&self) -> &BufferOptions {
144-
self.0.get_options()
145-
}
146-
fn obj_bytes(&self) -> BorrowedValue<[u8]> {
147-
self.0.obj_bytes()
148-
}
149-
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
150-
self.0.obj_bytes_mut()
151-
}
152-
fn release(&self) {}
153-
fn as_contiguous(&self) -> Option<BorrowedValue<[u8]>> {
154-
self.0.as_contiguous()
155-
}
156-
fn as_contiguous_mut(&self) -> Option<BorrowedValueMut<[u8]>> {
157-
self.0.as_contiguous_mut()
158-
}
159-
fn to_contiguous(&self) -> Vec<u8> {
160-
self.0.to_contiguous()
121+
impl Clone for PyBuffer {
122+
fn clone(&self) -> Self {
123+
self.clone_with_options(self.options.clone())
161124
}
162125
}
163126

‎vm/src/builtins/bytearray.rs

Lines changed: 13 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ use super::pystr::PyStrRef;
66
use super::pytype::PyTypeRef;
77
use super::tuple::PyTupleRef;
88
use crate::anystr::{self, AnyStr};
9-
use crate::buffer::{BufferOptions, PyBuffer, ResizeGuard};
9+
use crate::buffer::{BufferOptions, PyBuffer, PyBufferInternal, ResizeGuard};
1010
use crate::bytesinner::{
1111
bytes_decode, bytes_from_object, value_from_object, ByteInnerFindOptions, ByteInnerNewOptions,
1212
ByteInnerPaddingOptions, ByteInnerSplitOptions, ByteInnerTranslateOptions, DecodeArgs,
@@ -670,41 +670,34 @@ impl Comparable for PyByteArray {
670670
}
671671

672672
impl AsBuffer for PyByteArray {
673-
fn get_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<Box<dyn PyBuffer>> {
674-
zelf.exports.fetch_add(1);
675-
let buf = ByteArrayBuffer {
676-
bytearray: zelf.clone(),
677-
options: BufferOptions {
673+
fn get_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
674+
let buffer = PyBuffer::new(
675+
zelf.clone(),
676+
BufferOptions {
678677
readonly: false,
679678
len: zelf.len(),
680679
..Default::default()
681680
},
682-
};
683-
Ok(Box::new(buf))
681+
);
682+
Ok(buffer)
684683
}
685684
}
686685

687-
#[derive(Debug)]
688-
struct ByteArrayBuffer {
689-
bytearray: PyByteArrayRef,
690-
options: BufferOptions,
691-
}
692-
693-
impl PyBuffer for ByteArrayBuffer {
686+
impl PyBufferInternal for PyRef<PyByteArray> {
694687
fn obj_bytes(&self) -> BorrowedValue<[u8]> {
695-
self.bytearray.borrow_buf().into()
688+
self.borrow_buf().into()
696689
}
697690

698691
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
699-
PyRwLockWriteGuard::map(self.bytearray.inner_mut(), |inner| &mut *inner.elements).into()
692+
PyRwLockWriteGuard::map(self.inner_mut(), |inner| &mut *inner.elements).into()
700693
}
701694

702695
fn release(&self) {
703-
self.bytearray.exports.fetch_sub(1);
696+
self.exports.fetch_sub(1);
704697
}
705698

706-
fn get_options(&self) -> &BufferOptions {
707-
&self.options
699+
fn retain(&self) {
700+
self.exports.fetch_add(1);
708701
}
709702
}
710703

‎vm/src/builtins/bytes.rs

Lines changed: 10 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use super::int::PyIntRef;
33
use super::pystr::PyStrRef;
44
use super::pytype::PyTypeRef;
55
use crate::anystr::{self, AnyStr};
6-
use crate::buffer::{BufferOptions, PyBuffer};
6+
use crate::buffer::{BufferOptions, PyBuffer, PyBufferInternal};
77
use crate::builtins::tuple::PyTupleRef;
88
use crate::bytesinner::{
99
bytes_decode, ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions,
@@ -519,38 +519,29 @@ impl PyBytes {
519519
}
520520

521521
impl AsBuffer for PyBytes {
522-
fn get_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<Box<dyn PyBuffer>> {
523-
let buf = BytesBuffer {
524-
bytes: zelf.clone(),
525-
options: BufferOptions {
522+
fn get_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
523+
let buf = PyBuffer::new(
524+
zelf.clone(),
525+
BufferOptions {
526526
len: zelf.len(),
527527
..Default::default()
528528
},
529-
};
530-
Ok(Box::new(buf))
529+
);
530+
Ok(buf)
531531
}
532532
}
533533

534-
#[derive(Debug)]
535-
struct BytesBuffer {
536-
bytes: PyBytesRef,
537-
options: BufferOptions,
538-
}
539-
540-
impl PyBuffer for BytesBuffer {
534+
impl PyBufferInternal for PyRef<PyBytes> {
541535
fn obj_bytes(&self) -> BorrowedValue<[u8]> {
542-
self.bytes.as_bytes().into()
536+
self.as_bytes().into()
543537
}
544538

545539
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
546540
unreachable!("bytes is not mutable")
547541
}
548542

549543
fn release(&self) {}
550-
551-
fn get_options(&self) -> &BufferOptions {
552-
&self.options
553-
}
544+
fn retain(&self) {}
554545
}
555546

556547
impl Hashable for PyBytes {

0 commit comments

Comments
 (0)