Skip to content

Commit

Permalink
always use a Python iterator for sets and frozensets (#3849)
Browse files Browse the repository at this point in the history 
* always use a Python iterator for sets and frozensets

* add newsfragment
  • Loading branch information
@davidhewitt
davidhewitt committed Feb 17, 2024
1 parent 65cf580 commit eb90b81
Show file tree
Hide file tree
Showing 4 changed files with 62 additions and 184 deletions.
1 change: 1 addition & 0 deletions newsfragments/3849.added.md
View file
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
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
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
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));
});
}
}

0 comments on commit eb90b81

Please sign in to comment.