update memoryview eq

qingshi163 · qingshi163 · commit 896590c27d3e · 2020-10-08T12:17:54.000-05:00
diff --git a/vm/src/obj/objmemory.rs b/vm/src/obj/objmemory.rs
@@ -559,17 +559,7 @@ impl PyMemoryView {
         let bytes = &*zelf.obj_bytes();
         let elements: Vec<PyObjectRef> = (0..zelf.options.len)
             .map(|i| zelf.get_pos(i as isize).unwrap())
-            .map(|i| {
-                zelf.format_spec
-                    .unpack(&bytes[i..i + zelf.options.itemsize], vm)
-                    .map(|x| {
-                        if x.len() == 1 {
-                            x.fast_getitem(0)
-                        } else {
-                            x.into_object()
-                        }
-                    })
-            })
+            .map(|i| format_unpack(&zelf.format_spec, &bytes[i..i + zelf.options.itemsize], vm))
             .try_collect()?;
 
         Ok(PyList::from(elements).into_ref(vm))
@@ -635,13 +625,15 @@ impl PyMemoryView {
             return Ok(false);
         }
 
-        let options_cmp = |a: &BufferOptions, b: &BufferOptions| -> bool {
-            a.len == b.len && a.itemsize == b.itemsize
-        };
-
         let other = try_buffer_from_object(vm, other)?;
 
-        if !options_cmp(&zelf.options, &other.get_options()) {
+        let a_options = &zelf.options;
+        let b_options = &*other.get_options();
+
+        if a_options.len != b_options.len
+            || a_options.ndim != b_options.ndim
+            || a_options.shape != b_options.shape
+        {
             return Ok(false);
         }
 
@@ -670,7 +662,14 @@ impl PyMemoryView {
             }
         };
 
-        Ok(a == b)
+        if a_options.format == b_options.format {
+            Ok(a == b)
+        } else {
+            let a_list = unpack_bytes_seq_to_list(a, &a_options.format, vm)?;
+            let b_list = unpack_bytes_seq_to_list(b, &b_options.format, vm)?;
+
+            Ok(vm.bool_eq(a_list.as_object(), b_list.as_object())?)
+        }
     }
 }
 
@@ -789,3 +788,38 @@ pub fn try_buffer_from_object(vm: &VirtualMachine, obj: &PyObjectRef) -> PyResul
         obj_cls.name
     )))
 }
+
+fn format_unpack(
+    format_spec: &FormatSpec,
+    bytes: &[u8],
+    vm: &VirtualMachine,
+) -> PyResult<PyObjectRef> {
+    format_spec.unpack(bytes, vm).map(|x| {
+        if x.len() == 1 {
+            x.fast_getitem(0)
+        } else {
+            x.into_object()
+        }
+    })
+}
+
+pub fn unpack_bytes_seq_to_list(
+    bytes: &[u8],
+    format: &str,
+    vm: &VirtualMachine,
+) -> PyResult<PyListRef> {
+    let format_spec = PyMemoryView::parse_format(format, vm)?;
+    let itemsize = format_spec.size();
+
+    if bytes.len() % itemsize != 0 {
+        return Err(vm.new_value_error("bytes length not a multiple of item size".to_owned()));
+    }
+
+    let len = bytes.len() / itemsize;
+
+    let elements: Vec<PyObjectRef> = (0..len)
+        .map(|i| format_unpack(&format_spec, &bytes[i..i + itemsize], vm))
+        .try_collect()?;
+
+    Ok(PyList::from(elements).into_ref(vm))
+}
diff --git a/vm/src/stdlib/pystruct.rs b/vm/src/stdlib/pystruct.rs
@@ -193,7 +193,7 @@ pub(crate) mod _struct {
             Ok(PyTupleRef::with_elements(items, &vm.ctx))
         }
 
-        fn size(&self) -> usize {
+        pub fn size(&self) -> usize {
             self.codes.iter().map(FormatCode::size).sum()
         }
     }

Original file line number	Diff line number	Diff line change
`@@ -193,7 +193,7 @@ pub(crate) mod _struct {`
`193`	`193`	`Ok(PyTupleRef::with_elements(items, &vm.ctx))`
`194`	`194`	`}`
`195`	`195`
`196`		`- fn size(&self) -> usize {`
	`196`	`+ pub fn size(&self) -> usize {`
`197`	`197`	`self.codes.iter().map(FormatCode::size).sum()`
`198`	`198`	`}`
`199`	`199`	`}`