From eb90b81d4481767a49a04242785f1f2f46b007d6 Mon Sep 17 00:00:00 2001 From: David Hewitt Date: Sat, 17 Feb 2024 10:57:53 +0000 Subject: [PATCH] always use a Python iterator for sets and frozensets (#3849) * always use a Python iterator for sets and frozensets * add newsfragment --- newsfragments/3849.added.md | 1 + newsfragments/3849.changed.md | 1 + src/types/frozenset.rs | 117 ++++++++---------------------- src/types/set.rs | 131 ++++++++-------------------------- 4 files changed, 64 insertions(+), 186 deletions(-) create mode 100644 newsfragments/3849.added.md create mode 100644 newsfragments/3849.changed.md diff --git a/newsfragments/3849.added.md b/newsfragments/3849.added.md new file mode 100644 index 00000000..8aa9a55d --- /dev/null +++ b/newsfragments/3849.added.md @@ -0,0 +1 @@ +Implement `ExactSizeIterator` for `set` and `frozenset` iterators on `abi3` feature. diff --git a/newsfragments/3849.changed.md b/newsfragments/3849.changed.md new file mode 100644 index 00000000..ee8dfbf5 --- /dev/null +++ b/newsfragments/3849.changed.md @@ -0,0 +1 @@ +`PySet` and `PyFrozenSet` iterators now always iterate the equivalent of `iter(set)`. (A "fast path" with no noticeable performance benefit was removed.) diff --git a/src/types/frozenset.rs b/src/types/frozenset.rs index fa56a9c6..3be3f0ed 100644 --- a/src/types/frozenset.rs +++ b/src/types/frozenset.rs @@ -1,4 +1,3 @@ -#[cfg(Py_LIMITED_API)] use crate::types::PyIterator; use crate::{ err::{self, PyErr, PyResult}, @@ -205,6 +204,13 @@ impl<'py> Iterator for PyFrozenSetIterator<'py> { } } +impl ExactSizeIterator for PyFrozenSetIterator<'_> { + #[inline] + fn len(&self) -> usize { + self.0.len() + } +} + impl<'py> IntoIterator for &'py PyFrozenSet { type Item = &'py PyAny; type IntoIter = PyFrozenSetIterator<'py>; @@ -224,89 +230,41 @@ impl<'py> IntoIterator for Bound<'py, PyFrozenSet> { } } -#[cfg(Py_LIMITED_API)] -mod impl_ { - use super::*; +/// PyO3 implementation of an iterator for a Python `frozenset` object. +pub struct BoundFrozenSetIterator<'p> { + it: Bound<'p, PyIterator>, + // Remaining elements in the frozenset + remaining: usize, +} - /// PyO3 implementation of an iterator for a Python `set` object. - pub struct BoundFrozenSetIterator<'p> { - it: Bound<'p, PyIterator>, - } - - impl<'py> BoundFrozenSetIterator<'py> { - pub(super) fn new(frozenset: Bound<'py, PyFrozenSet>) -> Self { - Self { - it: PyIterator::from_bound_object(&frozenset).unwrap(), - } - } - } - - impl<'py> Iterator for BoundFrozenSetIterator<'py> { - type Item = Bound<'py, super::PyAny>; - - /// Advances the iterator and returns the next value. - #[inline] - fn next(&mut self) -> Option { - self.it.next().map(Result::unwrap) +impl<'py> BoundFrozenSetIterator<'py> { + pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self { + Self { + it: PyIterator::from_bound_object(&set).unwrap(), + remaining: set.len(), } } } -#[cfg(not(Py_LIMITED_API))] -mod impl_ { - use super::*; +impl<'py> Iterator for BoundFrozenSetIterator<'py> { + type Item = Bound<'py, super::PyAny>; - /// PyO3 implementation of an iterator for a Python `frozenset` object. - pub struct BoundFrozenSetIterator<'py> { - set: Bound<'py, PyFrozenSet>, - pos: ffi::Py_ssize_t, + /// Advances the iterator and returns the next value. + fn next(&mut self) -> Option { + self.remaining = self.remaining.saturating_sub(1); + self.it.next().map(Result::unwrap) } - impl<'py> BoundFrozenSetIterator<'py> { - pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self { - Self { set, pos: 0 } - } - } - - impl<'py> Iterator for BoundFrozenSetIterator<'py> { - type Item = Bound<'py, PyAny>; - - #[inline] - fn next(&mut self) -> Option { - unsafe { - let mut key: *mut ffi::PyObject = std::ptr::null_mut(); - let mut hash: ffi::Py_hash_t = 0; - if ffi::_PySet_NextEntry(self.set.as_ptr(), &mut self.pos, &mut key, &mut hash) != 0 - { - // _PySet_NextEntry returns borrowed object; for safety must make owned (see #890) - Some(key.assume_borrowed(self.set.py()).to_owned()) - } else { - None - } - } - } - - #[inline] - fn size_hint(&self) -> (usize, Option) { - let len = self.len(); - (len, Some(len)) - } - } - - impl<'py> ExactSizeIterator for BoundFrozenSetIterator<'py> { - fn len(&self) -> usize { - self.set.len().saturating_sub(self.pos as usize) - } - } - - impl<'py> ExactSizeIterator for PyFrozenSetIterator<'py> { - fn len(&self) -> usize { - self.0.len() - } + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) } } -pub use impl_::*; +impl<'py> ExactSizeIterator for BoundFrozenSetIterator<'py> { + fn len(&self) -> usize { + self.remaining + } +} #[inline] pub(crate) fn new_from_iter( @@ -387,7 +345,6 @@ mod tests { } #[test] - #[cfg(not(Py_LIMITED_API))] fn test_frozenset_iter_size_hint() { Python::with_gil(|py| { let set = PyFrozenSet::new(py, &[1]).unwrap(); @@ -402,18 +359,6 @@ mod tests { }); } - #[test] - #[cfg(Py_LIMITED_API)] - fn test_frozenset_iter_size_hint() { - Python::with_gil(|py| { - let set = PyFrozenSet::new(py, &[1]).unwrap(); - let iter = set.iter(); - - // No known bounds - assert_eq!(iter.size_hint(), (0, None)); - }); - } - #[test] fn test_frozenset_builder() { use super::PyFrozenSetBuilder; diff --git a/src/types/set.rs b/src/types/set.rs index 3204e71f..0c173352 100644 --- a/src/types/set.rs +++ b/src/types/set.rs @@ -1,4 +1,3 @@ -#[cfg(Py_LIMITED_API)] use crate::types::PyIterator; use crate::{ err::{self, PyErr, PyResult}, @@ -277,6 +276,12 @@ impl<'py> Iterator for PySetIterator<'py> { } } +impl ExactSizeIterator for PySetIterator<'_> { + fn len(&self) -> usize { + self.0.len() + } +} + impl<'py> IntoIterator for &'py PySet { type Item = &'py PyAny; type IntoIter = PySetIterator<'py>; @@ -304,103 +309,42 @@ impl<'py> IntoIterator for Bound<'py, PySet> { } } -#[cfg(Py_LIMITED_API)] -mod impl_ { - use super::*; +/// PyO3 implementation of an iterator for a Python `set` object. +pub struct BoundSetIterator<'p> { + it: Bound<'p, PyIterator>, + // Remaining elements in the set. This is fine to store because + // Python will error if the set changes size during iteration. + remaining: usize, +} - /// PyO3 implementation of an iterator for a Python `set` object. - pub struct BoundSetIterator<'p> { - it: Bound<'p, PyIterator>, - } - - impl<'py> BoundSetIterator<'py> { - pub(super) fn new(set: Bound<'py, PySet>) -> Self { - Self { - it: PyIterator::from_bound_object(&set).unwrap(), - } - } - } - - impl<'py> Iterator for BoundSetIterator<'py> { - type Item = Bound<'py, super::PyAny>; - - /// Advances the iterator and returns the next value. - /// - /// # Panics - /// - /// If PyO3 detects that the set is mutated during iteration, it will panic. - #[inline] - fn next(&mut self) -> Option { - self.it.next().map(Result::unwrap) +impl<'py> BoundSetIterator<'py> { + pub(super) fn new(set: Bound<'py, PySet>) -> Self { + Self { + it: PyIterator::from_bound_object(&set).unwrap(), + remaining: set.len(), } } } -#[cfg(not(Py_LIMITED_API))] -mod impl_ { - use super::*; +impl<'py> Iterator for BoundSetIterator<'py> { + type Item = Bound<'py, super::PyAny>; - /// PyO3 implementation of an iterator for a Python `set` object. - pub struct BoundSetIterator<'py> { - set: Bound<'py, super::PySet>, - pos: ffi::Py_ssize_t, - used: ffi::Py_ssize_t, + /// Advances the iterator and returns the next value. + fn next(&mut self) -> Option { + self.remaining = self.remaining.saturating_sub(1); + self.it.next().map(Result::unwrap) } - impl<'py> BoundSetIterator<'py> { - pub(super) fn new(set: Bound<'py, PySet>) -> Self { - let used = unsafe { ffi::PySet_Size(set.as_ptr()) }; - BoundSetIterator { set, pos: 0, used } - } - } - - impl<'py> Iterator for BoundSetIterator<'py> { - type Item = Bound<'py, super::PyAny>; - - /// Advances the iterator and returns the next value. - /// - /// # Panics - /// - /// If PyO3 detects that the set is mutated during iteration, it will panic. - #[inline] - fn next(&mut self) -> Option { - unsafe { - let len = ffi::PySet_Size(self.set.as_ptr()); - assert_eq!(self.used, len, "Set changed size during iteration"); - - let mut key: *mut ffi::PyObject = std::ptr::null_mut(); - let mut hash: ffi::Py_hash_t = 0; - if ffi::_PySet_NextEntry(self.set.as_ptr(), &mut self.pos, &mut key, &mut hash) != 0 - { - // _PySet_NextEntry returns borrowed object; for safety must make owned (see #890) - Some(key.assume_borrowed(self.set.py()).to_owned()) - } else { - None - } - } - } - - #[inline] - fn size_hint(&self) -> (usize, Option) { - let len = self.len(); - (len, Some(len)) - } - } - - impl<'py> ExactSizeIterator for BoundSetIterator<'py> { - fn len(&self) -> usize { - self.set.len().saturating_sub(self.pos as usize) - } - } - - impl<'py> ExactSizeIterator for PySetIterator<'py> { - fn len(&self) -> usize { - self.0.len() - } + fn size_hint(&self) -> (usize, Option) { + (self.remaining, Some(self.remaining)) } } -pub use impl_::*; +impl<'py> ExactSizeIterator for BoundSetIterator<'py> { + fn len(&self) -> usize { + self.remaining + } +} #[inline] pub(crate) fn new_from_iter( @@ -571,7 +515,6 @@ mod tests { } #[test] - #[cfg(not(Py_LIMITED_API))] fn test_set_iter_size_hint() { Python::with_gil(|py| { let set = PySet::new(py, &[1]).unwrap(); @@ -585,16 +528,4 @@ mod tests { assert_eq!(iter.size_hint(), (0, Some(0))); }); } - - #[test] - #[cfg(Py_LIMITED_API)] - fn test_set_iter_size_hint() { - Python::with_gil(|py| { - let set = PySet::new(py, &[1]).unwrap(); - let iter = set.iter(); - - // No known bounds - assert_eq!(iter.size_hint(), (0, None)); - }); - } }