Skip to content

Commit 88b32ec

Browse files
qingshi163qingshi163
committed
Fix io and array the right way to use buffer protocol
1 parent 52cb6fa commit 88b32ec

File tree

7 files changed

+97
-113
lines changed

7 files changed

+97
-113
lines changed

‎stdlib/src/array.rs

Lines changed: 15 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ pub(crate) use array::make_module;
33
#[pymodule(name = "array")]
44
mod array {
55
use crate::common::{
6-
borrow::{BorrowedValue, BorrowedValueMut},
76
lock::{
87
PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyRwLock, PyRwLockReadGuard,
98
PyRwLockWriteGuard,
@@ -20,7 +19,7 @@ mod array {
2019
ArgBytesLike, ArgIntoFloat, ArgIterable, IntoPyObject, IntoPyResult, OptionalArg,
2120
},
2221
protocol::{
23-
BufferInternal, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn,
22+
BufferMethods, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn,
2423
PyMappingMethods,
2524
},
2625
sliceable::{PySliceableSequence, PySliceableSequenceMut, SaturatedSlice, SequenceIndex},
@@ -1218,36 +1217,32 @@ mod array {
12181217
let array = zelf.read();
12191218
let buf = PyBuffer::new(
12201219
zelf.as_object().clone(),
1221-
PyArrayBufferInternal(zelf.clone()),
12221220
BufferOptions {
12231221
readonly: false,
12241222
len: array.len(),
12251223
itemsize: array.itemsize(),
12261224
format: array.typecode_str().into(),
12271225
..Default::default()
12281226
},
1227+
&BUFFER_METHODS,
12291228
);
12301229
Ok(buf)
12311230
}
12321231
}
12331232

1234-
impl BufferInternal for PyArray {
1235-
fn obj_bytes(&self) -> BorrowedValue<[u8]> {
1236-
self.get_bytes().into()
1237-
}
1238-
1239-
fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
1240-
self.get_bytes_mut().into()
1241-
}
1242-
1243-
fn release(&self) {
1244-
self.exports.fetch_sub(1);
1245-
}
1246-
1247-
fn retain(&self) {
1248-
self.exports.fetch_add(1);
1249-
}
1250-
}
1233+
static BUFFER_METHODS: BufferMethods = BufferMethods {
1234+
obj_bytes: |zelf| zelf.payload::<PyArray>().unwrap().get_bytes().into(),
1235+
obj_bytes_mut: |zelf| zelf.payload::<PyArray>().unwrap().get_bytes_mut().into(),
1236+
contiguous: None,
1237+
contiguous_mut: None,
1238+
collect_bytes: None,
1239+
release: Some(|zelf| {
1240+
zelf.payload::<PyArray>().unwrap().exports.fetch_sub(1);
1241+
}),
1242+
retain: Some(|zelf| {
1243+
zelf.payload::<PyArray>().unwrap().exports.fetch_add(1);
1244+
}),
1245+
};
12511246

12521247
impl AsMapping for PyArray {
12531248
fn as_mapping(_zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyMappingMethods {

‎vm/src/builtins/bytearray.rs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ use super::{
44
PyTypeRef,
55
};
66
use crate::common::{
7-
borrow::{BorrowedValue, BorrowedValueMut},
87
lock::{
98
PyMappedRwLockReadGuard, PyMappedRwLockWriteGuard, PyMutex, PyRwLock, PyRwLockReadGuard,
109
PyRwLockWriteGuard,
@@ -34,7 +33,6 @@ use crate::{
3433
};
3534
use bstr::ByteSlice;
3635
use crossbeam_utils::atomic::AtomicCell;
37-
use rustpython_common::atomic::Radium;
3836
use std::mem::size_of;
3937

4038
/// "bytearray(iterable_of_ints) -> bytearray\n\

‎vm/src/builtins/bytes.rs

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,7 @@ use crate::{
2020
TryFromBorrowedObject, TypeProtocol, VirtualMachine,
2121
};
2222
use bstr::ByteSlice;
23-
use rustpython_common::{
24-
borrow::{BorrowedValue, BorrowedValueMut},
25-
lock::PyMutex,
26-
};
23+
use rustpython_common::lock::PyMutex;
2724
use std::mem::size_of;
2825
use std::ops::Deref;
2926

‎vm/src/builtins/memory.rs

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,7 @@ impl PyMemoryView {
131131
}
132132

133133
#[pymethod]
134-
fn release(&self) {
134+
pub fn release(&self) {
135135
if self.released.compare_exchange(false, true).is_ok() {
136136
self.buffer._release();
137137
}
@@ -660,8 +660,6 @@ impl AsBuffer for PyMemoryView {
660660
}
661661
}
662662

663-
struct Released;
664-
665663
impl AsMapping for PyMemoryView {
666664
fn as_mapping(_zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyMappingMethods {
667665
PyMappingMethods {

‎vm/src/protocol/buffer.rs

Lines changed: 8 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -6,17 +6,17 @@ use std::{borrow::Cow, fmt::Debug};
66

77
pub struct BufferMethods {
88
// always reflecting the whole bytes of the most top object
9-
obj_bytes: fn(&PyObjectRef) -> BorrowedValue<[u8]>,
9+
pub obj_bytes: fn(&PyObjectRef) -> BorrowedValue<[u8]>,
1010
// always reflecting the whole bytes of the most top object
11-
obj_bytes_mut: fn(&PyObjectRef) -> BorrowedValueMut<[u8]>,
11+
pub obj_bytes_mut: fn(&PyObjectRef) -> BorrowedValueMut<[u8]>,
1212
// GUARANTEE: called only if the buffer option is contiguous
13-
contiguous: Option<fn(&PyObjectRef) -> BorrowedValue<[u8]>>,
13+
pub contiguous: Option<fn(&PyObjectRef) -> BorrowedValue<[u8]>>,
1414
// GUARANTEE: called only if the buffer option is contiguous
15-
contiguous_mut: Option<fn(&PyObjectRef) -> BorrowedValueMut<[u8]>>,
15+
pub contiguous_mut: Option<fn(&PyObjectRef) -> BorrowedValueMut<[u8]>>,
1616
// collect bytes to buf when buffer option is not contiguous
17-
collect_bytes: Option<fn(&PyObjectRef, buf: &mut Vec<u8>)>,
18-
release: Option<fn(&PyObjectRef)>,
19-
retain: Option<fn(&PyObjectRef)>,
17+
pub collect_bytes: Option<fn(&PyObjectRef, buf: &mut Vec<u8>)>,
18+
pub release: Option<fn(&PyObjectRef)>,
19+
pub retain: Option<fn(&PyObjectRef)>,
2020
}
2121

2222
impl Debug for BufferMethods {
@@ -68,7 +68,7 @@ impl PyBuffer {
6868

6969
pub fn contiguous_or_collect<R, F: FnOnce(&[u8]) -> R>(&self, f: F) -> R {
7070
let borrowed;
71-
let collected;
71+
let mut collected;
7272
let v = if self.options.contiguous {
7373
borrowed = self._contiguous();
7474
&*borrowed
@@ -84,10 +84,6 @@ impl PyBuffer {
8484
Self::new(self.obj.clone(), options, self.methods)
8585
}
8686

87-
pub fn move_with_options(self, options: BufferOptions) -> Self {
88-
Self { options, ..self }
89-
}
90-
9187
// SAFETY: should only called if option has contiguous
9288
pub(crate) fn _contiguous(&self) -> BorrowedValue<[u8]> {
9389
self.methods

‎vm/src/stdlib/io.rs

Lines changed: 66 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,6 @@ mod _io {
7474
PyMappedThreadMutexGuard, PyMutex, PyRwLock, PyRwLockReadGuard, PyRwLockWriteGuard,
7575
PyThreadMutex, PyThreadMutexGuard,
7676
},
77-
rc::PyRc,
7877
};
7978
use crate::{
8079
builtins::{
@@ -86,10 +85,10 @@ mod _io {
8685
OptionalOption,
8786
},
8887
protocol::{BufferMethods, BufferOptions, BufferResizeGuard, PyBuffer, PyIterReturn},
89-
types::{Constructor, Destructor, IterNext, Iterable},
88+
types::{AsBuffer, Constructor, Destructor, IterNext, Iterable},
9089
utils::Either,
9190
vm::{ReprGuard, VirtualMachine},
92-
IdProtocol, PyContext, PyObjectRef, PyRef, PyResult, PyValue, StaticType,
91+
IdProtocol, PyContext, PyObjectRef, PyObjectWrap, PyRef, PyResult, PyValue, StaticType,
9392
TryFromBorrowedObject, TryFromObject, TypeProtocol,
9493
};
9594
use bstr::ByteSlice;
@@ -882,16 +881,17 @@ mod _io {
882881
};
883882
// TODO: see if we can encapsulate this pattern in a function in memory.rs like
884883
// fn slice_as_memory<R>(s: &[u8], f: impl FnOnce(PyMemoryViewRef) -> R) -> R
885-
let writebuf = PyRc::new(BufferedRawBuffer {
884+
let writebuf = BufferedRawBuffer {
886885
data: std::mem::take(&mut self.buffer).into(),
887886
range: buf_range,
888-
});
887+
}
888+
.into_ref(vm);
889889
let raw = self.raw.as_ref().unwrap();
890890
let memobj = PyMemoryView::from_buffer(
891891
PyBuffer {
892-
obj: raw.clone(),
892+
obj: writebuf.clone().into_object(),
893893
options,
894-
internal: writebuf.clone(),
894+
methods: &BUFFERED_RAW_BUFFER_METHODS,
895895
},
896896
vm,
897897
)?
@@ -1121,15 +1121,16 @@ mod _io {
11211121
};
11221122
// TODO: see if we can encapsulate this pattern in a function in memory.rs like
11231123
// fn slice_as_memory<R>(s: &[u8], f: impl FnOnce(PyMemoryViewRef) -> R) -> R
1124-
let readbuf = PyRc::new(BufferedRawBuffer {
1124+
let readbuf = BufferedRawBuffer {
11251125
data: std::mem::take(v).into(),
11261126
range: buf_range,
1127-
});
1127+
}
1128+
.into_ref(vm);
11281129
let memobj = PyMemoryView::from_buffer(
11291130
PyBuffer {
1130-
obj: vm.ctx.none(),
1131+
obj: readbuf.clone().into_object(),
11311132
options,
1132-
internal: readbuf.clone(),
1133+
methods: &BUFFERED_RAW_BUFFER_METHODS,
11331134
},
11341135
vm,
11351136
)?
@@ -1323,37 +1324,50 @@ mod _io {
13231324
// the memoryobj for the buffer until after the BufferedIO is destroyed, you
13241325
// can get a use-after-free, so this is a bit safe
13251326
#[pyclass(noattr, module = false, name = "_buffered_raw_buffer")]
1326-
#[derive(Debug)]
1327+
#[derive(Debug, PyValue)]
13271328
struct BufferedRawBuffer {
13281329
data: PyMutex<Vec<u8>>,
13291330
range: Range<usize>,
13301331
}
1331-
#[pyimpl(flags(BASETYPE))]
1332+
#[pyimpl(with(AsBuffer))]
13321333
impl BufferedRawBuffer {}
13331334

13341335
static BUFFERED_RAW_BUFFER_METHODS: BufferMethods = BufferMethods {
13351336
obj_bytes: |zelf| {
1336-
zelf.downcast_ref::<BufferedRawBuffer>(zelf)
1337-
.unwrap()
1338-
.data
1339-
.lock()
1340-
.into()
1337+
let zelf = zelf.downcast_ref::<BufferedRawBuffer>().unwrap();
1338+
BorrowedValue::map(zelf.data.lock().into(), |x| x.as_slice())
13411339
},
1340+
obj_bytes_mut: |zelf| {
1341+
let zelf = zelf.downcast_ref::<BufferedRawBuffer>().unwrap();
1342+
BorrowedValueMut::map(zelf.data.lock().into(), |x| x.as_mut_slice())
1343+
},
1344+
contiguous: Some(|zelf| {
1345+
let zelf = zelf.downcast_ref::<BufferedRawBuffer>().unwrap();
1346+
BorrowedValue::map(zelf.data.lock().into(), |x| &x[zelf.range.clone()])
1347+
}),
1348+
contiguous_mut: Some(|zelf| {
1349+
let zelf = zelf.downcast_ref::<BufferedRawBuffer>().unwrap();
1350+
BorrowedValueMut::map(zelf.data.lock().into(), |x| &mut x[zelf.range.clone()])
1351+
}),
1352+
collect_bytes: None,
1353+
// TODO: export counting to make sure no exports before retrive data to buffer
1354+
release: None,
1355+
retain: None,
13421356
};
1343-
// impl BufferInternal for BufferedRawBuffer {
1344-
// fn obj_bytes(&self) -> BorrowedValue<[u8]> {
1345-
// BorrowedValue::map(self.data.lock().into(), |data| &data[self.range.clone()])
1346-
// }
1347-
1348-
// fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
1349-
// BorrowedValueMut::map(self.data.lock().into(), |data| {
1350-
// &mut data[self.range.clone()]
1351-
// })
1352-
// }
13531357

1354-
// fn release(&self) {}
1355-
// fn retain(&self) {}
1356-
// }
1358+
impl AsBuffer for BufferedRawBuffer {
1359+
fn as_buffer(zelf: &PyRef<Self>, _vm: &VirtualMachine) -> PyResult<PyBuffer> {
1360+
let buf = PyBuffer::new(
1361+
zelf.as_object().clone(),
1362+
BufferOptions {
1363+
len: zelf.range.end - zelf.range.start,
1364+
..Default::default()
1365+
},
1366+
&BUFFERED_RAW_BUFFER_METHODS,
1367+
);
1368+
Ok(buf)
1369+
}
1370+
}
13571371

13581372
pub fn get_offset(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult<Offset> {
13591373
use std::convert::TryInto;
@@ -3345,30 +3359,32 @@ mod _io {
33453359
len: self.buffer.read().cursor.get_ref().len(),
33463360
..Default::default()
33473361
};
3348-
let buffer = PyBuffer::new(self.as_object().clone(), self, options);
3362+
let buffer = PyBuffer::new(self.as_object().clone(), options, &BYTES_IO_BUFFER_METHODS);
33493363
let view = PyMemoryView::from_buffer(buffer, vm)?;
33503364
Ok(view)
33513365
}
33523366
}
33533367

3354-
// impl BufferInternal for PyRef<BytesIO> {
3355-
// fn obj_bytes(&self) -> BorrowedValue<[u8]> {
3356-
// PyRwLockReadGuard::map(self.buffer.read(), |x| x.cursor.get_ref().as_slice()).into()
3357-
// }
3358-
3359-
// fn obj_bytes_mut(&self) -> BorrowedValueMut<[u8]> {
3360-
// PyRwLockWriteGuard::map(self.buffer.write(), |x| x.cursor.get_mut().as_mut_slice())
3361-
// .into()
3362-
// }
3363-
3364-
// fn release(&self) {
3365-
// self.exports.fetch_sub(1);
3366-
// }
3367-
3368-
// fn retain(&self) {
3369-
// self.exports.fetch_add(1);
3370-
// }
3371-
// }
3368+
static BYTES_IO_BUFFER_METHODS: BufferMethods = BufferMethods {
3369+
obj_bytes: |zelf| {
3370+
let zelf = zelf.downcast_ref::<BytesIO>().unwrap();
3371+
PyRwLockReadGuard::map(zelf.buffer.read(), |x| x.cursor.get_ref().as_slice()).into()
3372+
},
3373+
obj_bytes_mut: |zelf| {
3374+
let zelf = zelf.downcast_ref::<BytesIO>().unwrap();
3375+
PyRwLockWriteGuard::map(zelf.buffer.write(), |x| x.cursor.get_mut().as_mut_slice())
3376+
.into()
3377+
},
3378+
contiguous: None,
3379+
contiguous_mut: None,
3380+
collect_bytes: None,
3381+
release: Some(|zelf| {
3382+
zelf.downcast_ref::<BytesIO>().unwrap().exports.fetch_sub(1);
3383+
}),
3384+
retain: Some(|zelf| {
3385+
zelf.downcast_ref::<BytesIO>().unwrap().exports.fetch_add(1);
3386+
}),
3387+
};
33723388

33733389
impl<'a> BufferResizeGuard<'a> for BytesIO {
33743390
type Resizable = PyRwLockWriteGuard<'a, BufferedIO>;

‎vm/src/stdlib/sre.rs

Lines changed: 6 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -140,31 +140,15 @@ mod _sre {
140140
vm: &VirtualMachine,
141141
f: F,
142142
) -> PyResult<R> {
143-
let buffer;
144-
let guard;
145-
let vec;
146-
let s;
147-
let str_drive = if self.isbytes {
148-
buffer = PyBuffer::try_from_borrowed_object(vm, &string)?;
149-
let bytes = match buffer.as_contiguous() {
150-
Some(bytes) => {
151-
guard = bytes;
152-
&*guard
153-
}
154-
None => {
155-
vec = buffer.to_contiguous();
156-
vec.as_slice()
157-
}
158-
};
159-
StrDrive::Bytes(bytes)
143+
if self.isbytes {
144+
let buffer = PyBuffer::try_from_borrowed_object(vm, &string)?;
145+
buffer.contiguous_or_collect(|bytes| f(StrDrive::Bytes(bytes)))
160146
} else {
161-
s = string
147+
let s = string
162148
.payload::<PyStr>()
163149
.ok_or_else(|| vm.new_type_error("expected string".to_owned()))?;
164-
StrDrive::Str(s.as_str())
165-
};
166-
167-
f(str_drive)
150+
f(StrDrive::Str(s.as_str()))
151+
}
168152
}
169153

170154
fn with_state<R, F: FnOnce(State) -> PyResult<R>>(

0 commit comments

Comments
 (0)