Skip to content

Commit df78720

Browse files
committed
memoryview eq
1 parent 39e5543 commit df78720

File tree

3 files changed

+81
-52
lines changed

3 files changed

+81
-52
lines changed

vm/src/obj/objmemory.rs

Lines changed: 79 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,19 @@
11
use std::{fmt::Debug, iter::FromIterator};
22

3-
use super::{
3+
use crate::obj::objbytes::{PyBytes, PyBytesRef};
4+
use crate::obj::objlist::{PyList, PyListRef};
5+
use crate::obj::{
46
objsequence::get_pos, objsequence::get_slice_range, objsequence::SequenceIndex,
57
objslice::PySliceRef, objstr::PyString, objtype::PyClassRef,
68
};
7-
use crate::common::hash::PyHash;
8-
use crate::obj::objbytes::{PyBytes, PyBytesRef};
9-
use crate::obj::objlist::{PyList, PyListRef};
10-
use crate::pyobject::IdProtocol;
119
use crate::pyobject::{
12-
PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyThreadingConstraint, PyValue,
13-
TypeProtocol,
10+
IdProtocol, PyClassImpl, PyContext, PyObjectRef, PyRef, PyResult, PyThreadingConstraint,
11+
PyValue, TypeProtocol,
1412
};
1513
use crate::slots::{Bufferable, Hashable};
1614
use crate::VirtualMachine;
1715
use crate::{bytesinner::try_as_bytes, pyobject::IntoPyObject};
16+
use crate::{common::hash::PyHash, pyobject::PyComparisonValue};
1817
use crossbeam_utils::atomic::AtomicCell;
1918
use num_bigint::BigInt;
2019
use num_traits::{One, Signed, ToPrimitive, Zero};
@@ -40,7 +39,6 @@ pub struct BufferOptions {
4039
pub len: usize,
4140
pub itemsize: usize,
4241
pub contiguous: bool,
43-
// pub range: Range<usize>,
4442
// TODO: support multiple dimension array
4543
pub ndim: usize,
4644
// TODO: support format
@@ -142,17 +140,20 @@ impl PyMemoryView {
142140
.map(|_| PyString::from(&self.options.format))
143141
}
144142

145-
fn getitem_by_idx(&self, i: isize, vm: &VirtualMachine) -> PyResult {
143+
fn get_pos(&self, i: isize) -> Option<usize> {
146144
let len = self.options.len;
147-
let i =
148-
get_pos(i, len).ok_or_else(|| vm.new_index_error("index out of range".to_owned()))?;
149-
// let i = (self.start as isize + i as isize * self.step) as usize;
150-
let i = if self.step < 0 {
145+
let i = get_pos(i, len)?;
146+
Some(if self.step < 0 {
151147
self.stop - (-self.step) as usize * i + self.step as usize
152148
} else {
153149
self.start + self.step as usize * i
154-
};
150+
})
151+
}
155152

153+
fn getitem_by_idx(&self, i: isize, vm: &VirtualMachine) -> PyResult {
154+
let i = self
155+
.get_pos(i)
156+
.ok_or_else(|| vm.new_index_error("index out of range".to_owned()))?;
156157
unpack_value(self.buffer.as_bytes()[i], &self.options.format, vm)
157158
}
158159

@@ -198,30 +199,10 @@ impl PyMemoryView {
198199
let range = get_slice_range(&start, &stop, len);
199200
let range = zelf.start + range.start..zelf.start + range.end;
200201

201-
// let step: BigInt = step * self.step;
202-
// newlen will be 0 if step is overflowed, in that case step should not change the behavour
202+
// newlen will be 0 if step is overflowed
203203
let newstep = newstep.to_isize().unwrap_or(-1);
204204

205-
// let (start, stop, step, newlen) = if range.start >= range.end {
206-
// (0, 0, 1, 0)
207-
// } else if step.is_positive() {
208-
// let start = self.start + range.start;
209-
// match step.to_isize() {
210-
// Some(step) => {
211-
// let newlen = len / step as usize;
212-
// (start, start + newlen, step, newlen)
213-
// },
214-
// None => (start, start + 1, 1, 1),
215-
// }
216-
// } else {
217-
// match step.to_isize() {
218-
// Some(step) => (range.end - 1, step, len / step.abs() as usize),
219-
// None => (range.end - 1, -1, 1),
220-
// }
221-
// };
222-
223205
let obj = zelf.obj.clone();
224-
// let buffer = try_buffer_from_object(self.obj.clone(), vm)?;
225206
let buffer = Box::new(zelf.clone());
226207
let options = zelf.options.clone();
227208

@@ -244,7 +225,6 @@ impl PyMemoryView {
244225
#[pymethod(magic)]
245226
fn getitem(zelf: PyRef<Self>, needle: SequenceIndex, vm: &VirtualMachine) -> PyResult {
246227
zelf.try_not_released(vm)?;
247-
// self.obj.get_item(needle, vm)
248228
match needle {
249229
SequenceIndex::Int(i) => zelf.getitem_by_idx(i, vm),
250230
SequenceIndex::Slice(slice) => Self::getitem_by_slice(zelf, slice, vm),
@@ -256,17 +236,41 @@ impl PyMemoryView {
256236
self.try_not_released(vm).map(|_| self.options.len)
257237
}
258238

239+
fn get_contiguous_bytes(&self) -> Option<BorrowedValue<[u8]>> {
240+
if self.options.contiguous {
241+
Some(BorrowedValue::map(self.buffer.as_bytes(), |x| {
242+
&x[self.start..self.stop]
243+
}))
244+
} else {
245+
None
246+
}
247+
}
248+
259249
#[pymethod]
260250
fn tobytes(&self, vm: &VirtualMachine) -> PyResult<PyBytesRef> {
261-
self.try_not_released(vm)
262-
.map(|_| PyBytes::from(self.buffer.as_bytes().to_vec()).into_ref(vm))
251+
self.try_not_released(vm)?;
252+
if let Some(bytes) = self.get_contiguous_bytes() {
253+
Ok(PyBytes::from(bytes.to_vec()).into_ref(vm))
254+
} else {
255+
let bytes = &*self.buffer.as_bytes();
256+
let bytes = (0..self.options.len)
257+
.map(|i| self.get_pos(i as isize).unwrap())
258+
.map(|i| bytes[i])
259+
.collect::<Vec<u8>>();
260+
Ok(PyBytes::from(bytes).into_ref(vm))
261+
}
263262
}
264263

265264
#[pymethod]
266265
fn tolist(&self, vm: &VirtualMachine) -> PyResult<PyListRef> {
267266
self.try_not_released(vm).map(|_| {
268-
PyList::from_iter(self.buffer.as_bytes().iter().map(|x| x.into_pyobject(vm)))
269-
.into_ref(vm)
267+
let bytes = &*self.buffer.as_bytes();
268+
PyList::from_iter(
269+
(0..self.options.len)
270+
.map(|i| self.get_pos(i as isize).unwrap())
271+
.map(|i| bytes[i].into_pyobject(vm)),
272+
)
273+
.into_ref(vm)
270274
})
271275
}
272276

@@ -296,6 +300,39 @@ impl PyMemoryView {
296300
format!("<memory at 0x{:x}>", zelf.get_id())
297301
}
298302
}
303+
304+
#[pymethod(magic)]
305+
fn eq(
306+
zelf: PyRef<Self>,
307+
other: PyObjectRef,
308+
vm: &VirtualMachine,
309+
) -> PyResult<PyComparisonValue> {
310+
zelf.try_not_released(vm)?;
311+
if zelf.is(&other) {
312+
return Ok(PyComparisonValue::Implemented(true));
313+
}
314+
let options_cmp = |a: &BufferOptions, b: &BufferOptions| -> bool {
315+
a.len == b.len && a.itemsize == b.itemsize && a.format == b.format
316+
};
317+
if let Some(other) = other.payload_if_subclass::<PyMemoryView>(vm) {
318+
if options_cmp(&zelf.options, &other.options) {
319+
let a = zelf.tolist(vm)?;
320+
let b = other.tolist(vm)?;
321+
if vm.bool_eq(a.into_object(), b.into_object())? {
322+
return Ok(PyComparisonValue::Implemented(true));
323+
}
324+
}
325+
return Ok(PyComparisonValue::Implemented(false));
326+
}
327+
if let Ok(buffer) = try_buffer_from_object(other, vm) {
328+
let options = buffer.get_options();
329+
if options_cmp(&zelf.options, &options) && (**(zelf.tobytes(vm)?) == *buffer.as_bytes())
330+
{
331+
return Ok(PyComparisonValue::Implemented(true));
332+
}
333+
}
334+
Ok(PyComparisonValue::Implemented(false))
335+
}
299336
}
300337

301338
impl Drop for PyMemoryView {
@@ -305,14 +342,13 @@ impl Drop for PyMemoryView {
305342
}
306343

307344
impl Bufferable for PyMemoryView {
308-
fn get_buffer(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult<Box<dyn BufferProtocol>> {
309-
try_buffer_from_object(zelf.obj.clone(), vm)
345+
fn get_buffer(zelf: PyRef<Self>, _vm: &VirtualMachine) -> PyResult<Box<dyn BufferProtocol>> {
346+
Ok(Box::new(zelf))
310347
}
311348
}
312349

313350
impl BufferProtocol for PyMemoryViewRef {
314351
fn get_options(&self) -> BufferOptions {
315-
// use options from ourself
316352
self.options.clone()
317353
}
318354

vm/src/stdlib/array.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -834,7 +834,7 @@ impl Bufferable for PyArray {
834834
}
835835

836836
impl BufferProtocol for PyArrayRef {
837-
fn as_bytes(&self) -> BorrowedValue<[u8]>{
837+
fn as_bytes(&self) -> BorrowedValue<[u8]> {
838838
self.get_bytes().into()
839839
}
840840

vm/src/stdlib/marshal.rs

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,27 +7,20 @@ mod decl {
77
use crate::pyobject::PyResult;
88
use crate::vm::VirtualMachine;
99
use crate::{bytecode, obj::objmemory::try_buffer_from_object, pyobject::PyObjectRef};
10-
use crate::pyobject::TypeProtocol;
1110

1211
#[pyfunction]
1312
fn dumps(co: PyCodeRef) -> PyBytes {
1413
PyBytes::from(co.code.to_bytes())
1514
}
1615

1716
#[pyfunction]
18-
// fn loads(code_bytes: PyBytesRef, vm: &VirtualMachine) -> PyResult<PyCode> {
1917
fn loads(code_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyCode> {
20-
dbg!(&code_bytes.class().name);
2118
let buffer = try_buffer_from_object(code_bytes, vm)?;
2219
let options = buffer.get_options();
23-
dbg!(&options);
2420
if !options.contiguous {
2521
return Err(vm.new_value_error("buffer is not contiguous".to_owned()));
2622
}
27-
let bytes = buffer.as_bytes().to_vec();
28-
// dbg!(&*buffer.as_bytes());
29-
// let code = bytecode::CodeObject::from_bytes(&code_bytes)
30-
let code = bytecode::CodeObject::from_bytes(bytes.as_slice())
23+
let code = bytecode::CodeObject::from_bytes(&*buffer.as_bytes())
3124
.map_err(|_| vm.new_value_error("Couldn't deserialize python bytecode".to_owned()))?;
3225
Ok(PyCode { code })
3326
}

0 commit comments

Comments
 (0)