From da3e37e335bfe108662c96330330a6dd4c53a181 Mon Sep 17 00:00:00 2001 From: Bas Schoenmaeckers <7943856+bschoenmaeckers@users.noreply.github.com> Date: Wed, 19 Feb 2025 18:06:39 +0100 Subject: [PATCH] Optimize `Iterator::last()` for `PyList` & `PyTuple` and `Iterator::count()` for `PyDict`, `PyList`, `PyTuple` & `PySet` (#4878) * Optimize `Iterator::last()` for `PyList` & `PyTuple` * Optimize `Iterator::count()` for `PyDict`, `PyList`, `PyTuple` & `PySet` --- newsfragments/4878.added.md | 2 ++ src/types/dict.rs | 24 ++++++++++++++++++ src/types/frozenset.rs | 16 ++++++++++++ src/types/list.rs | 33 +++++++++++++++++++++++++ src/types/set.rs | 16 ++++++++++++ src/types/tuple.rs | 49 +++++++++++++++++++++++++++++++++++++ 6 files changed, 140 insertions(+) create mode 100644 newsfragments/4878.added.md diff --git a/newsfragments/4878.added.md b/newsfragments/4878.added.md new file mode 100644 index 00000000000..0130b2b805b --- /dev/null +++ b/newsfragments/4878.added.md @@ -0,0 +1,2 @@ +- Optimizes `last` for `BoundListIterator`, `BoundTupleIterator` and `BorrowedTupleIterator` +- Optimizes `Iterator::count()` for `PyDict`, `PyList`, `PyTuple` & `PySet` \ No newline at end of file diff --git a/src/types/dict.rs b/src/types/dict.rs index 0d2e6ff335f..3f5118aab6b 100644 --- a/src/types/dict.rs +++ b/src/types/dict.rs @@ -534,6 +534,14 @@ impl<'py> Iterator for BoundDictIterator<'py> { (len, Some(len)) } + #[inline] + fn count(self) -> usize + where + Self: Sized, + { + self.len() + } + #[inline] #[cfg(Py_GIL_DISABLED)] fn fold(mut self, init: B, mut f: F) -> B @@ -736,6 +744,14 @@ mod borrowed_iter { let len = self.len(); (len, Some(len)) } + + #[inline] + fn count(self) -> usize + where + Self: Sized, + { + self.len() + } } impl ExactSizeIterator for BorrowedDictIter<'_, '_> { @@ -1657,4 +1673,12 @@ mod tests { .is_err()); }); } + + #[test] + fn test_iter_count() { + Python::with_gil(|py| { + let dict = [(1, 1), (2, 2), (3, 3)].into_py_dict(py).unwrap(); + assert_eq!(dict.iter().count(), 3); + }) + } } diff --git a/src/types/frozenset.rs b/src/types/frozenset.rs index 954c49b5902..df22560cd8e 100644 --- a/src/types/frozenset.rs +++ b/src/types/frozenset.rs @@ -237,6 +237,14 @@ impl<'py> Iterator for BoundFrozenSetIterator<'py> { fn size_hint(&self) -> (usize, Option) { (self.remaining, Some(self.remaining)) } + + #[inline] + fn count(self) -> usize + where + Self: Sized, + { + self.len() + } } impl ExactSizeIterator for BoundFrozenSetIterator<'_> { @@ -358,4 +366,12 @@ mod tests { assert!(!set.contains(3).unwrap()); }); } + + #[test] + fn test_iter_count() { + Python::with_gil(|py| { + let set = PyFrozenSet::new(py, vec![1, 2, 3]).unwrap(); + assert_eq!(set.iter().count(), 3); + }) + } } diff --git a/src/types/list.rs b/src/types/list.rs index e16e5009a01..ead22315f05 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -695,6 +695,22 @@ impl<'py> Iterator for BoundListIterator<'py> { (len, Some(len)) } + #[inline] + fn count(self) -> usize + where + Self: Sized, + { + self.len() + } + + #[inline] + fn last(mut self) -> Option + where + Self: Sized, + { + self.next_back() + } + #[inline] #[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))] fn fold(mut self, init: B, mut f: F) -> B @@ -1778,4 +1794,21 @@ mod tests { assert_eq!(iter4.next_back().unwrap().extract::().unwrap(), 5); }) } + + #[test] + fn test_iter_last() { + Python::with_gil(|py| { + let list = PyList::new(py, vec![1, 2, 3]).unwrap(); + let last = list.iter().last(); + assert_eq!(last.unwrap().extract::().unwrap(), 3); + }) + } + + #[test] + fn test_iter_count() { + Python::with_gil(|py| { + let list = PyList::new(py, vec![1, 2, 3]).unwrap(); + assert_eq!(list.iter().count(), 3); + }) + } } diff --git a/src/types/set.rs b/src/types/set.rs index 60aa9428562..ddaddbfe5ed 100644 --- a/src/types/set.rs +++ b/src/types/set.rs @@ -267,6 +267,14 @@ impl<'py> Iterator for BoundSetIterator<'py> { fn size_hint(&self) -> (usize, Option) { (self.remaining, Some(self.remaining)) } + + #[inline] + fn count(self) -> usize + where + Self: Sized, + { + self.len() + } } impl ExactSizeIterator for BoundSetIterator<'_> { @@ -479,4 +487,12 @@ mod tests { assert_eq!(iter.size_hint(), (0, Some(0))); }); } + + #[test] + fn test_iter_count() { + Python::with_gil(|py| { + let set = PySet::new(py, vec![1, 2, 3]).unwrap(); + assert_eq!(set.iter().count(), 3); + }) + } } diff --git a/src/types/tuple.rs b/src/types/tuple.rs index 1d674f47a79..81a7ad911e9 100644 --- a/src/types/tuple.rs +++ b/src/types/tuple.rs @@ -377,6 +377,22 @@ impl<'py> Iterator for BoundTupleIterator<'py> { (len, Some(len)) } + #[inline] + fn count(self) -> usize + where + Self: Sized, + { + self.len() + } + + #[inline] + fn last(mut self) -> Option + where + Self: Sized, + { + self.next_back() + } + #[inline] #[cfg(not(feature = "nightly"))] fn nth(&mut self, n: usize) -> Option { @@ -548,6 +564,22 @@ impl<'a, 'py> Iterator for BorrowedTupleIterator<'a, 'py> { let len = self.len(); (len, Some(len)) } + + #[inline] + fn count(self) -> usize + where + Self: Sized, + { + self.len() + } + + #[inline] + fn last(mut self) -> Option + where + Self: Sized, + { + self.next_back() + } } impl DoubleEndedIterator for BorrowedTupleIterator<'_, '_> { @@ -1728,4 +1760,21 @@ mod tests { assert_eq!(iter4.next_back().unwrap().extract::().unwrap(), 5); }) } + + #[test] + fn test_iter_last() { + Python::with_gil(|py| { + let tuple = PyTuple::new(py, vec![1, 2, 3]).unwrap(); + let last = tuple.iter().last(); + assert_eq!(last.unwrap().extract::().unwrap(), 3); + }) + } + + #[test] + fn test_iter_count() { + Python::with_gil(|py| { + let tuple = PyTuple::new(py, vec![1, 2, 3]).unwrap(); + assert_eq!(tuple.iter().count(), 3); + }) + } }