Skip to content

Commit

Permalink
Implement nth, nth_back, advance_by and advance_back_by for tuple ite…
Browse files Browse the repository at this point in the history 
…rator
  • Loading branch information
@Owen-CH-Leung
Owen-CH-Leung committed Feb 9, 2025
1 parent 8ab39d2 commit 453d163
Show file tree
Hide file tree
Showing 2 changed files with 238 additions and 3 deletions.
28 changes: 28 additions & 0 deletions pyo3-benches/benches/bench_tuple.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,10 +132,38 @@ fn tuple_into_pyobject(b: &mut Bencher<'_>) {
});
}

fn tuple_nth(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
const LEN: usize = 50;
let list = PyTuple::new(py, 0..LEN).unwrap();
let mut sum = 0;
b.iter(|| {
for i in 0..LEN {
sum += list.iter().nth(i).unwrap().extract::<usize>().unwrap();
}
});
});
}

fn tuple_nth_back(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
const LEN: usize = 50;
let list = PyTuple::new(py, 0..LEN).unwrap();
let mut sum = 0;
b.iter(|| {
for i in 0..LEN {
sum += list.iter().nth_back(i).unwrap().extract::<usize>().unwrap();
}
});
});
}

fn criterion_benchmark(c: &mut Criterion) {
c.bench_function("iter_tuple", iter_tuple);
c.bench_function("tuple_new", tuple_new);
c.bench_function("tuple_get_item", tuple_get_item);
c.bench_function("tuple_nth", tuple_nth);
c.bench_function("tuple_nth_back", tuple_nth_back);
#[cfg(not(any(Py_LIMITED_API, PyPy)))]
c.bench_function("tuple_get_item_unchecked", tuple_get_item_unchecked);
c.bench_function("tuple_get_borrowed_item", tuple_get_borrowed_item);
Expand Down
213 changes: 210 additions & 3 deletions src/types/tuple.rs
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,3 @@
use std::iter::FusedIterator;

use crate::ffi::{self, Py_ssize_t};
use crate::ffi_ptr_ext::FfiPtrExt;
#[cfg(feature = "experimental-inspect")]
Expand All @@ -13,6 +11,9 @@ use crate::{
};
#[allow(deprecated)]
use crate::{IntoPy, ToPyObject};
use std::iter::FusedIterator;
#[cfg(feature = "nightly")]
use std::num::NonZero;

#[inline]
#[track_caller]
Expand Down Expand Up @@ -375,6 +376,48 @@ impl<'py> Iterator for BoundTupleIterator<'py> {
let len = self.len();
(len, Some(len))
}

#[inline]
#[cfg(not(feature = "nightly"))]
fn nth(&mut self, n: usize) -> Option<Self::Item> {
let length = self.length.min(self.tuple.len());
let target_index = self.index + n;
if target_index < length {
let item = {
{
unsafe { self.tuple.get_item_unchecked(target_index) }
}
};
self.index = target_index + 1;
Some(item)
} else {
None
}
}

#[inline]
#[cfg(feature = "nightly")]
fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
let max_len = self.length.min(self.tuple.len());
let currently_at = self.index;
if currently_at >= max_len {
if n == 0 {
return Ok(());
} else {
return Err(unsafe { NonZero::new_unchecked(n) });
}
}

let items_left = max_len - currently_at;
if n <= items_left {
self.index += n;
Ok(())
} else {
self.index = max_len;
let remainder = n - items_left;
Err(unsafe { NonZero::new_unchecked(remainder) })
}
}
}

impl DoubleEndedIterator for BoundTupleIterator<'_> {
Expand All @@ -391,6 +434,48 @@ impl DoubleEndedIterator for BoundTupleIterator<'_> {
None
}
}

#[inline]
#[cfg(not(feature = "nightly"))]
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
let length_size = self.length.min(self.tuple.len());
if self.index + n < length_size {
let target_index = length_size - n - 1;
let item = {
{
unsafe { self.tuple.get_item_unchecked(target_index) }
}
};
self.length = target_index;
Some(item)
} else {
None
}
}

#[inline]
#[cfg(feature = "nightly")]
fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
let max_len = self.length.min(self.tuple.len());
let currently_at = self.index;
if currently_at >= max_len {
if n == 0 {
return Ok(());
} else {
return Err(unsafe { NonZero::new_unchecked(n) });
}
}

let items_left = max_len - currently_at;
if n <= items_left {
self.length = max_len - n;
Ok(())
} else {
self.length = max_len;
let remainder = n - items_left;
Err(unsafe { NonZero::new_unchecked(remainder) })
}
}
}

impl ExactSizeIterator for BoundTupleIterator<'_> {
Expand Down Expand Up @@ -979,8 +1064,9 @@ mod tests {
use crate::types::{any::PyAnyMethods, tuple::PyTupleMethods, PyList, PyTuple};
use crate::{IntoPyObject, Python};
use std::collections::HashSet;
#[cfg(feature = "nightly")]
use std::num::NonZero;
use std::ops::Range;

#[test]
fn test_new() {
Python::with_gil(|py| {
Expand Down Expand Up @@ -1523,4 +1609,125 @@ mod tests {
}
})
}

#[test]
fn test_bound_tuple_nth() {
Python::with_gil(|py| {
let tuple = PyTuple::new(py, vec![1, 2, 3, 4]).unwrap();
let mut iter = tuple.iter();
assert_eq!(iter.nth(1).unwrap().extract::<i32>().unwrap(), 2);
assert_eq!(iter.nth(1).unwrap().extract::<i32>().unwrap(), 4);
assert!(iter.nth(1).is_none());

let tuple = PyTuple::new(py, Vec::<i32>::new()).unwrap();
let mut iter = tuple.iter();
iter.next();
assert!(iter.nth(1).is_none());

let tuple = PyTuple::new(py, vec![1, 2, 3]).unwrap();
let mut iter = tuple.iter();
assert!(iter.nth(10).is_none());

let tuple = PyTuple::new(py, vec![6, 7, 8, 9, 10]).unwrap();
let mut iter = tuple.iter();
assert_eq!(iter.next().unwrap().extract::<i32>().unwrap(), 6);
assert_eq!(iter.nth(2).unwrap().extract::<i32>().unwrap(), 9);
assert_eq!(iter.next().unwrap().extract::<i32>().unwrap(), 10);

let mut iter = tuple.iter();
assert_eq!(iter.nth_back(1).unwrap().extract::<i32>().unwrap(), 9);
assert_eq!(iter.nth(2).unwrap().extract::<i32>().unwrap(), 8);
assert!(iter.next().is_none());
});
}

#[test]
fn test_bound_tuple_nth_back() {
Python::with_gil(|py| {
let tuple = PyTuple::new(py, vec![1, 2, 3, 4, 5]).unwrap();
let mut iter = tuple.iter();
assert_eq!(iter.nth_back(0).unwrap().extract::<i32>().unwrap(), 5);
assert_eq!(iter.nth_back(1).unwrap().extract::<i32>().unwrap(), 3);
assert!(iter.nth_back(2).is_none());

let tuple = PyTuple::new(py, Vec::<i32>::new()).unwrap();
let mut iter = tuple.iter();
assert!(iter.nth_back(0).is_none());
assert!(iter.nth_back(1).is_none());

let tuple = PyTuple::new(py, vec![1, 2, 3]).unwrap();
let mut iter = tuple.iter();
assert!(iter.nth_back(5).is_none());

let tuple = PyTuple::new(py, vec![1, 2, 3, 4, 5]).unwrap();
let mut iter = tuple.iter();
iter.next_back(); // Consume the last element
assert_eq!(iter.nth_back(1).unwrap().extract::<i32>().unwrap(), 3);
assert_eq!(iter.next_back().unwrap().extract::<i32>().unwrap(), 2);
assert_eq!(iter.nth_back(0).unwrap().extract::<i32>().unwrap(), 1);

let tuple = PyTuple::new(py, vec![1, 2, 3, 4, 5]).unwrap();
let mut iter = tuple.iter();
assert_eq!(iter.nth_back(1).unwrap().extract::<i32>().unwrap(), 4);
assert_eq!(iter.nth_back(2).unwrap().extract::<i32>().unwrap(), 1);

let mut iter2 = tuple.iter();
iter2.next_back();
assert_eq!(iter2.nth_back(1).unwrap().extract::<i32>().unwrap(), 3);
assert_eq!(iter2.next_back().unwrap().extract::<i32>().unwrap(), 2);

let mut iter3 = tuple.iter();
iter3.nth(1);
assert_eq!(iter3.nth_back(2).unwrap().extract::<i32>().unwrap(), 3);
assert!(iter3.nth_back(0).is_none());
});
}

#[cfg(feature = "nightly")]
#[test]
fn test_bound_tuple_advance_by() {
Python::with_gil(|py| {
let tuple = PyTuple::new(py, vec![1, 2, 3, 4, 5]).unwrap();
let mut iter = tuple.iter();

assert_eq!(iter.advance_by(2), Ok(()));
assert_eq!(iter.next().unwrap().extract::<i32>().unwrap(), 3);
assert_eq!(iter.advance_by(0), Ok(()));
assert_eq!(iter.advance_by(100), Err(NonZero::new(98).unwrap()));

let mut iter2 = tuple.iter();
assert_eq!(iter2.advance_by(6), Err(NonZero::new(1).unwrap()));

let mut iter3 = tuple.iter();
assert_eq!(iter3.advance_by(5), Ok(()));

let mut iter4 = tuple.iter();
assert_eq!(iter4.advance_by(0), Ok(()));
assert_eq!(iter4.next().unwrap().extract::<i32>().unwrap(), 1);
})
}

#[cfg(feature = "nightly")]
#[test]
fn test_bound_tuple_advance_back_by() {
Python::with_gil(|py| {
let tuple = PyTuple::new(py, vec![1, 2, 3, 4, 5]).unwrap();
let mut iter = tuple.iter();

assert_eq!(iter.advance_back_by(2), Ok(()));
assert_eq!(iter.next_back().unwrap().extract::<i32>().unwrap(), 3);
assert_eq!(iter.advance_back_by(0), Ok(()));
assert_eq!(iter.advance_back_by(100), Err(NonZero::new(98).unwrap()));

let mut iter2 = tuple.iter();
assert_eq!(iter2.advance_back_by(6), Err(NonZero::new(1).unwrap()));

let mut iter3 = tuple.iter();
assert_eq!(iter3.advance_back_by(5), Ok(()));

let mut iter4 = tuple.iter();
assert_eq!(iter4.advance_back_by(0), Ok(()));
assert_eq!(iter4.next_back().unwrap().extract::<i32>().unwrap(), 5);
})
}
}

0 comments on commit 453d163

Please sign in to comment.