diff --git a/benches/microbenchmarks/list_count_builtins.py b/benches/microbenchmarks/list_count_builtins.py new file mode 100644 index 000000000..ac716c2cd --- /dev/null +++ b/benches/microbenchmarks/list_count_builtins.py @@ -0,0 +1,4 @@ +l = [i for i in range(ITERATIONS)] + +# --- +l.count(1) \ No newline at end of file diff --git a/benches/microbenchmarks/list_count_custom_eq.py b/benches/microbenchmarks/list_count_custom_eq.py new file mode 100644 index 000000000..9db832a90 --- /dev/null +++ b/benches/microbenchmarks/list_count_custom_eq.py @@ -0,0 +1,8 @@ +class A: + def __eq__(self, other): + return True + +l = [A()] * ITERATIONS + +# --- +l.count(1) \ No newline at end of file diff --git a/extra_tests/snippets/list.py b/extra_tests/snippets/list.py index d0f24000d..6ce9c0a8f 100644 --- a/extra_tests/snippets/list.py +++ b/extra_tests/snippets/list.py @@ -646,7 +646,7 @@ m = MutatingCompare() l = [1, 2, 3, m, 4] m.list = l -l.count(4) # TODO: assert l.count(4) == 1 +assert l.count(4) == 1 l = [1, 2, 3, m, 4] m.list = l diff --git a/vm/src/builtins/list.rs b/vm/src/builtins/list.rs index 5ed4f8620..8f020c709 100644 --- a/vm/src/builtins/list.rs +++ b/vm/src/builtins/list.rs @@ -6,21 +6,21 @@ use crate::{ function::{ArgIterable, FuncArgs, IntoPyObject, OptionalArg}, protocol::{PyIterReturn, PyMappingMethods}, sequence::{self, SimpleSeq}, - sliceable::{PySliceableSequence, PySliceableSequenceMut, SequenceIndex}, + sliceable::{saturate_index, PySliceableSequence, PySliceableSequenceMut, SequenceIndex}, stdlib::sys, types::{ - AsMapping, Comparable, Constructor, Hashable, IterNext, IterNextIterable, Iterable, - PyComparisonOp, Unconstructible, Unhashable, + richcompare_wrapper, AsMapping, Comparable, Constructor, Hashable, IterNext, + IterNextIterable, Iterable, PyComparisonOp, RichCompareFunc, Unconstructible, Unhashable, }, utils::Either, vm::{ReprGuard, VirtualMachine}, - PyClassDef, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, PyRef, PyResult, PyValue, - TryFromObject, TypeProtocol, + IdProtocol, PyClassDef, PyClassImpl, PyComparisonValue, PyContext, PyObjectRef, PyRef, + PyResult, PyValue, TryFromObject, TypeProtocol, }; use std::fmt; use std::iter::FromIterator; use std::mem::size_of; -use std::ops::DerefMut; +use std::ops::{DerefMut, Range}; /// Built-in mutable sequence. /// @@ -252,30 +252,173 @@ impl PyList { Ok(zelf.clone()) } + fn _iter_equal( + &self, + needle: &PyObjectRef, + range: Range, + mut f: F, + vm: &VirtualMachine, + ) -> PyResult { + let needle_cls = needle.class(); + let needle_cmp = needle_cls + .mro_find_map(|cls| cls.slots.richcompare.load()) + .unwrap(); + + let mut borrower = None; + let mut i = range.start; + + let index = loop { + if i >= range.end { + break usize::MAX; + } + let guard = if let Some(x) = borrower.take() { + x + } else { + self.borrow_vec() + }; + + let elem = if let Some(x) = guard.get(i) { + x + } else { + break usize::MAX; + }; + + if elem.is(needle) { + f(); + if SHORT { + break i; + } + borrower = Some(guard); + } else { + let elem_cls = elem.class(); + let reverse_first = !elem_cls.is(&needle_cls) && elem_cls.issubclass(&needle_cls); + + let eq = if reverse_first { + let elem_cmp = elem_cls + .mro_find_map(|cls| cls.slots.richcompare.load()) + .unwrap(); + drop(elem_cls); + + fn cmp( + elem: &PyObjectRef, + needle: &PyObjectRef, + elem_cmp: RichCompareFunc, + needle_cmp: RichCompareFunc, + vm: &VirtualMachine, + ) -> PyResult { + match elem_cmp(elem, needle, PyComparisonOp::Eq, vm)? { + Either::B(PyComparisonValue::Implemented(value)) => Ok(value), + Either::A(obj) if !obj.is(&vm.ctx.not_implemented) => { + obj.try_to_bool(vm) + } + _ => match needle_cmp(needle, elem, PyComparisonOp::Eq, vm)? { + Either::B(PyComparisonValue::Implemented(value)) => Ok(value), + Either::A(obj) if !obj.is(&vm.ctx.not_implemented) => { + obj.try_to_bool(vm) + } + _ => Ok(false), + }, + } + } + + if elem_cmp as usize == richcompare_wrapper as usize { + let elem = elem.clone(); + drop(guard); + cmp(&elem, needle, elem_cmp, needle_cmp, vm)? + } else { + let eq = cmp(elem, needle, elem_cmp, needle_cmp, vm)?; + borrower = Some(guard); + eq + } + } else { + match needle_cmp(needle, elem, PyComparisonOp::Eq, vm)? { + Either::B(PyComparisonValue::Implemented(value)) => { + drop(elem_cls); + borrower = Some(guard); + value + } + Either::A(obj) if !obj.is(&vm.ctx.not_implemented) => { + drop(elem_cls); + borrower = Some(guard); + obj.try_to_bool(vm)? + } + _ => { + let elem_cmp = elem_cls + .mro_find_map(|cls| cls.slots.richcompare.load()) + .unwrap(); + drop(elem_cls); + + fn cmp( + elem: &PyObjectRef, + needle: &PyObjectRef, + elem_cmp: RichCompareFunc, + vm: &VirtualMachine, + ) -> PyResult { + match elem_cmp(elem, needle, PyComparisonOp::Eq, vm)? { + Either::B(PyComparisonValue::Implemented(value)) => Ok(value), + Either::A(obj) if !obj.is(&vm.ctx.not_implemented) => { + obj.try_to_bool(vm) + } + _ => Ok(false), + } + } + + if elem_cmp as usize == richcompare_wrapper as usize { + let elem = elem.clone(); + drop(guard); + cmp(&elem, needle, elem_cmp, vm)? + } else { + let eq = cmp(elem, needle, elem_cmp, vm)?; + borrower = Some(guard); + eq + } + } + } + }; + + if eq { + f(); + if SHORT { + break i; + } + } + } + i += 1; + }; + + // TODO: Optioned + Ok(index) + } + + fn foreach_equal( + &self, + needle: &PyObjectRef, + f: F, + vm: &VirtualMachine, + ) -> PyResult<()> { + self._iter_equal::<_, false>(needle, 0..usize::MAX, f, vm) + .map(|_| ()) + } + + fn find_equal( + &self, + needle: &PyObjectRef, + range: Range, + vm: &VirtualMachine, + ) -> PyResult { + self._iter_equal::<_, true>(needle, range, || {}, vm) + } + #[pymethod] fn count(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { - // TODO: to_vec() cause copy which leads to cost O(N). It need to be improved. - let elements = self.borrow_vec().to_vec(); - let mut count: usize = 0; - for elem in elements.iter() { - if vm.identical_or_equal(elem, &needle)? { - count += 1; - } - } + let mut count = 0; + self.foreach_equal(&needle, || count += 1, vm)?; Ok(count) } #[pymethod(magic)] - pub fn contains(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { - // TODO: to_vec() cause copy which leads to cost O(N). It need to be improved. - let elements = self.borrow_vec().to_vec(); - for elem in elements.iter() { - if vm.identical_or_equal(elem, &needle)? { - return Ok(true); - } - } - - Ok(false) + pub(crate) fn contains(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { + Ok(self.find_equal(&needle, 0..usize::MAX, vm)? != usize::MAX) } #[pymethod] @@ -286,33 +429,17 @@ impl PyList { stop: OptionalArg, vm: &VirtualMachine, ) -> PyResult { - let mut start = start.into_option().unwrap_or(0); - if start < 0 { - start += self.borrow_vec().len() as isize; - if start < 0 { - start = 0; - } + let len = self.len(); + let start = start.map(|i| saturate_index(i, len)).unwrap_or(0); + let stop = stop + .map(|i| saturate_index(i, len)) + .unwrap_or(sys::MAXSIZE as usize); + let index = self.find_equal(&needle, start..stop, vm)?; + if index == usize::MAX { + Err(vm.new_value_error(format!("'{}' is not in list", vm.to_str(&needle)?))) + } else { + Ok(index) } - let mut stop = stop.into_option().unwrap_or(sys::MAXSIZE); - if stop < 0 { - stop += self.borrow_vec().len() as isize; - if stop < 0 { - stop = 0; - } - } - // TODO: to_vec() cause copy which leads to cost O(N). It need to be improved. - let elements = self.borrow_vec().to_vec(); - for (index, element) in elements - .iter() - .enumerate() - .take(stop as usize) - .skip(start as usize) - { - if vm.identical_or_equal(element, &needle)? { - return Ok(index); - } - } - Err(vm.new_value_error(format!("'{}' is not in list", vm.to_str(&needle)?))) } #[pymethod] @@ -333,17 +460,9 @@ impl PyList { #[pymethod] fn remove(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - // TODO: to_vec() cause copy which leads to cost O(N). It need to be improved. - let elements = self.borrow_vec().to_vec(); - let mut ri: Option = None; - for (index, element) in elements.iter().enumerate() { - if vm.identical_or_equal(element, &needle)? { - ri = Some(index); - break; - } - } + let index = self.find_equal(&needle, 0..usize::MAX, vm)?; - if let Some(index) = ri { + if index != usize::MAX { // defer delete out of borrow Ok(self.borrow_vec_mut().remove(index)) } else { diff --git a/vm/src/types/slot.rs b/vm/src/types/slot.rs index ea5a1449c..6157b2b84 100644 --- a/vm/src/types/slot.rs +++ b/vm/src/types/slot.rs @@ -225,7 +225,7 @@ fn setattro_wrapper( Ok(()) } -fn richcompare_wrapper( +pub(crate) fn richcompare_wrapper( zelf: &PyObjectRef, other: &PyObjectRef, op: PyComparisonOp,