Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

always use a Python iterator for sets and frozensets #3849

Merged
merged 2 commits into from
Feb 17, 2024
Merged

always use a Python iterator for sets and frozensets #3849

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e95050c
always use a Python iterator for sets and frozensets
davidhewitt Feb 17, 2024
11bc256
add newsfragment
davidhewitt Feb 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions newsfragments/3849.added.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Implement `ExactSizeIterator` for `set` and `frozenset` iterators on `abi3` feature.
1 change: 1 addition & 0 deletions newsfragments/3849.changed.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.)
115 changes: 30 additions & 85 deletions src/types/frozenset.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
#[cfg(Py_LIMITED_API)]
use crate::types::PyIterator;
use crate::{
err::{self, PyErr, PyResult},
Expand Down Expand Up @@ -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>;
Expand All @@ -224,90 +230,42 @@ impl<'py> IntoIterator for Bound<'py, PyFrozenSet> {
}
}

#[cfg(Py_LIMITED_API)]
mod impl_ {
use super::*;

/// 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::Item> {
self.it.next().map(Result::unwrap)
}
}
/// 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,
}

#[cfg(not(Py_LIMITED_API))]
mod impl_ {
use super::*;

/// PyO3 implementation of an iterator for a Python `frozenset` object.
pub struct BoundFrozenSetIterator<'py> {
set: Bound<'py, PyFrozenSet>,
pos: ffi::Py_ssize_t,
}

impl<'py> BoundFrozenSetIterator<'py> {
pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self {
Self { set, pos: 0 }
impl<'py> BoundFrozenSetIterator<'py> {
pub(super) fn new(set: Bound<'py, PyFrozenSet>) -> Self {
Self {
it: PyIterator::from_bound_object(&set).unwrap(),
remaining: set.len(),
}
}
}

impl<'py> Iterator for BoundFrozenSetIterator<'py> {
type Item = Bound<'py, PyAny>;

#[inline]
fn next(&mut self) -> Option<Self::Item> {
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
}
}
}
impl<'py> Iterator for BoundFrozenSetIterator<'py> {
type Item = Bound<'py, super::PyAny>;

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.len();
(len, Some(len))
}
/// Advances the iterator and returns the next value.
fn next(&mut self) -> Option<Self::Item> {
self.remaining = self.remaining.saturating_sub(1);
self.it.next().map(Result::unwrap)
}

impl<'py> ExactSizeIterator for BoundFrozenSetIterator<'py> {
fn len(&self) -> usize {
self.set.len().saturating_sub(self.pos as usize)
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.remaining, Some(self.remaining))
}
}

impl<'py> ExactSizeIterator for PyFrozenSetIterator<'py> {
fn len(&self) -> usize {
self.0.len()
}
impl<'py> ExactSizeIterator for BoundFrozenSetIterator<'py> {
fn len(&self) -> usize {
self.remaining
}
}

pub use impl_::*;

#[inline]
pub(crate) fn new_from_iter<T: ToPyObject>(
py: Python<'_>,
Expand Down Expand Up @@ -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();
Expand All @@ -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;
Expand Down
129 changes: 30 additions & 99 deletions src/types/set.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
#[cfg(Py_LIMITED_API)]
use crate::types::PyIterator;
use crate::{
err::{self, PyErr, PyResult},
Expand Down Expand Up @@ -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>;
Expand Down Expand Up @@ -304,104 +309,43 @@ 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>,
}

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::Item> {
self.it.next().map(Result::unwrap)
}
}
/// 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,
}

#[cfg(not(Py_LIMITED_API))]
mod impl_ {
use super::*;

/// 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,
}

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> BoundSetIterator<'py> {
pub(super) fn new(set: Bound<'py, PySet>) -> Self {
Self {
it: PyIterator::from_bound_object(&set).unwrap(),
remaining: set.len(),
}
}
}

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::Item> {
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
}
}
}
impl<'py> Iterator for BoundSetIterator<'py> {
type Item = Bound<'py, super::PyAny>;

#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.len();
(len, Some(len))
}
/// Advances the iterator and returns the next value.
fn next(&mut self) -> Option<Self::Item> {
self.remaining = self.remaining.saturating_sub(1);
self.it.next().map(Result::unwrap)
}

impl<'py> ExactSizeIterator for BoundSetIterator<'py> {
fn len(&self) -> usize {
self.set.len().saturating_sub(self.pos as usize)
}
fn size_hint(&self) -> (usize, Option<usize>) {
(self.remaining, Some(self.remaining))
}
}

impl<'py> ExactSizeIterator for PySetIterator<'py> {
fn len(&self) -> usize {
self.0.len()
}
impl<'py> ExactSizeIterator for BoundSetIterator<'py> {
fn len(&self) -> usize {
self.remaining
}
}

pub use impl_::*;

#[inline]
pub(crate) fn new_from_iter<T: ToPyObject>(
py: Python<'_>,
Expand Down Expand Up @@ -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();
Expand All @@ -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));
});
}
}
Loading