Fix equality check in list_count, list_index and list_contains

Author: calixteman

tests/snippets/list.py

@@ -92,3 +92,21 @@
 x = [1]
 x.append(x)
 assert x in x
+assert x.index(x) == 1
+assert x.count(x) == 1
+
+class Foo(object):
+    def __eq__(self, x):
+        return False
+
+foo = Foo()
+foo1 = Foo()
+x = [1, foo, 2, foo, []]
+assert foo in x
+assert 2 in x
+assert x.index(foo) == 1
+assert x.count(foo) == 2
+assert x.index(2) == 2
+assert [] in x
+assert x.index([]) == 4
+assert foo1 not in x

vm/src/obj/objlist.rs

@@ -1,5 +1,5 @@
 use super::super::pyobject::{
-    PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
+    IdProtocol, PyContext, PyFuncArgs, PyObject, PyObjectPayload, PyObjectRef, PyResult, TypeProtocol,
 };
 use super::super::vm::{ReprGuard, VirtualMachine};
 use super::objbool;
@@ -234,9 +234,13 @@ fn list_count(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
     let elements = get_elements(zelf);
     let mut count: usize = 0;
     for element in elements.iter() {
-        let is_eq = vm._eq(element.clone(), value.clone())?;
-        if objbool::boolval(vm, is_eq)? {
+        if value.is(&element) {
             count += 1;
+        } else {
+            let is_eq = vm._eq(element.clone(), value.clone())?;
+            if objbool::boolval(vm, is_eq)? {
+                count += 1;
+            }
         }
     }
     Ok(vm.context().new_int(count))
@@ -262,6 +266,9 @@ fn list_index(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(list, Some(vm.ctx.list_type())), (needle, None)]
     );
     for (index, element) in get_elements(list).iter().enumerate() {
+        if needle.is(&element) {
+            return Ok(vm.context().new_int(index));
+        }
         let py_equal = vm._eq(needle.clone(), element.clone())?;
         if objbool::get_value(&py_equal) {
             return Ok(vm.context().new_int(index));
@@ -335,6 +342,9 @@ fn list_contains(vm: &mut VirtualMachine, args: PyFuncArgs) -> PyResult {
         required = [(list, Some(vm.ctx.list_type())), (needle, None)]
     );
     for element in get_elements(list).iter() {
+        if needle.is(&element) {
+            return Ok(vm.new_bool(true));
+        }
         match vm._eq(needle.clone(), element.clone()) {
             Ok(value) => {
                 if objbool::get_value(&value) {

vm/src/vm.rs

@@ -608,13 +608,9 @@ impl VirtualMachine {
     }
 
     pub fn _eq(&mut self, a: PyObjectRef, b: PyObjectRef) -> PyResult {
-        if a.is(&b) {
-            Ok(self.new_bool(true))
-        } else {
-            self.call_or_unsupported(a, b, "__eq__", "__eq__", |vm, _, _| {
-                Ok(vm.new_bool(false))
-            })
-        }
+        self.call_or_unsupported(a, b, "__eq__", "__eq__", |vm, a, b| {
+            Ok(vm.new_bool(a.is(&b)))
+        })
     }
 
     pub fn _ne(&mut self, a: PyObjectRef, b: PyObjectRef) -> PyResult {