Skip to content

Commit 6028cfc

Browse files
Owen-CH-LeungOwen-CH-Leung
authored
Optimize nth and nth_back for BoundTupleIterator (#4897)
* Implement nth, nth_back, advance_by and advance_back_by for tuple iterator * Add newsfragment * Fix failing clippy * Fix incorrect advance_back_by logic for tuple and list
1 parent 2c732a7 commit 6028cfc

File tree

4 files changed

+238
-4
lines changed

4 files changed

+238
-4
lines changed

‎newsfragments/4897.added.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Optimizes `nth`, `nth_back`, `advance_by` and `advance_back_by` for `BoundTupleIterator`

‎pyo3-benches/benches/bench_tuple.rs

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,10 +132,38 @@ fn tuple_into_pyobject(b: &mut Bencher<'_>) {
132132
});
133133
}
134134

135+
fn tuple_nth(b: &mut Bencher<'_>) {
136+
Python::with_gil(|py| {
137+
const LEN: usize = 50;
138+
let list = PyTuple::new(py, 0..LEN).unwrap();
139+
let mut sum = 0;
140+
b.iter(|| {
141+
for i in 0..LEN {
142+
sum += list.iter().nth(i).unwrap().extract::<usize>().unwrap();
143+
}
144+
});
145+
});
146+
}
147+
148+
fn tuple_nth_back(b: &mut Bencher<'_>) {
149+
Python::with_gil(|py| {
150+
const LEN: usize = 50;
151+
let list = PyTuple::new(py, 0..LEN).unwrap();
152+
let mut sum = 0;
153+
b.iter(|| {
154+
for i in 0..LEN {
155+
sum += list.iter().nth_back(i).unwrap().extract::<usize>().unwrap();
156+
}
157+
});
158+
});
159+
}
160+
135161
fn criterion_benchmark(c: &mut Criterion) {
136162
c.bench_function("iter_tuple", iter_tuple);
137163
c.bench_function("tuple_new", tuple_new);
138164
c.bench_function("tuple_get_item", tuple_get_item);
165+
c.bench_function("tuple_nth", tuple_nth);
166+
c.bench_function("tuple_nth_back", tuple_nth_back);
139167
#[cfg(not(any(Py_LIMITED_API, PyPy)))]
140168
c.bench_function("tuple_get_item_unchecked", tuple_get_item_unchecked);
141169
c.bench_function("tuple_get_borrowed_item", tuple_get_borrowed_item);

‎src/types/list.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -920,7 +920,7 @@ impl DoubleEndedIterator for BoundListIterator<'_> {
920920
length.0 = max_len - n;
921921
Ok(())
922922
} else {
923-
length.0 = max_len;
923+
length.0 = currently_at;
924924
let remainder = n - items_left;
925925
Err(unsafe { NonZero::new_unchecked(remainder) })
926926
}

‎src/types/tuple.rs

Lines changed: 208 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
use std::iter::FusedIterator;
2-
31
use crate::ffi::{self, Py_ssize_t};
42
use crate::ffi_ptr_ext::FfiPtrExt;
53
#[cfg(feature = "experimental-inspect")]
@@ -13,6 +11,9 @@ use crate::{
1311
};
1412
#[allow(deprecated)]
1513
use crate::{IntoPy, ToPyObject};
14+
use std::iter::FusedIterator;
15+
#[cfg(feature = "nightly")]
16+
use std::num::NonZero;
1617

1718
#[inline]
1819
#[track_caller]
@@ -375,6 +376,46 @@ impl<'py> Iterator for BoundTupleIterator<'py> {
375376
let len = self.len();
376377
(len, Some(len))
377378
}
379+
380+
#[inline]
381+
#[cfg(not(feature = "nightly"))]
382+
fn nth(&mut self, n: usize) -> Option<Self::Item> {
383+
let length = self.length.min(self.tuple.len());
384+
let target_index = self.index + n;
385+
if target_index < length {
386+
let item = unsafe {
387+
BorrowedTupleIterator::get_item(self.tuple.as_borrowed(), target_index).to_owned()
388+
};
389+
self.index = target_index + 1;
390+
Some(item)
391+
} else {
392+
None
393+
}
394+
}
395+
396+
#[inline]
397+
#[cfg(feature = "nightly")]
398+
fn advance_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
399+
let max_len = self.length.min(self.tuple.len());
400+
let currently_at = self.index;
401+
if currently_at >= max_len {
402+
if n == 0 {
403+
return Ok(());
404+
} else {
405+
return Err(unsafe { NonZero::new_unchecked(n) });
406+
}
407+
}
408+
409+
let items_left = max_len - currently_at;
410+
if n <= items_left {
411+
self.index += n;
412+
Ok(())
413+
} else {
414+
self.index = max_len;
415+
let remainder = n - items_left;
416+
Err(unsafe { NonZero::new_unchecked(remainder) })
417+
}
418+
}
378419
}
379420

380421
impl DoubleEndedIterator for BoundTupleIterator<'_> {
@@ -391,6 +432,46 @@ impl DoubleEndedIterator for BoundTupleIterator<'_> {
391432
None
392433
}
393434
}
435+
436+
#[inline]
437+
#[cfg(not(feature = "nightly"))]
438+
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
439+
let length_size = self.length.min(self.tuple.len());
440+
if self.index + n < length_size {
441+
let target_index = length_size - n - 1;
442+
let item = unsafe {
443+
BorrowedTupleIterator::get_item(self.tuple.as_borrowed(), target_index).to_owned()
444+
};
445+
self.length = target_index;
446+
Some(item)
447+
} else {
448+
None
449+
}
450+
}
451+
452+
#[inline]
453+
#[cfg(feature = "nightly")]
454+
fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero<usize>> {
455+
let max_len = self.length.min(self.tuple.len());
456+
let currently_at = self.index;
457+
if currently_at >= max_len {
458+
if n == 0 {
459+
return Ok(());
460+
} else {
461+
return Err(unsafe { NonZero::new_unchecked(n) });
462+
}
463+
}
464+
465+
let items_left = max_len - currently_at;
466+
if n <= items_left {
467+
self.length = max_len - n;
468+
Ok(())
469+
} else {
470+
self.length = currently_at;
471+
let remainder = n - items_left;
472+
Err(unsafe { NonZero::new_unchecked(remainder) })
473+
}
474+
}
394475
}
395476

396477
impl ExactSizeIterator for BoundTupleIterator<'_> {
@@ -979,8 +1060,9 @@ mod tests {
9791060
use crate::types::{any::PyAnyMethods, tuple::PyTupleMethods, PyList, PyTuple};
9801061
use crate::{IntoPyObject, Python};
9811062
use std::collections::HashSet;
1063+
#[cfg(feature = "nightly")]
1064+
use std::num::NonZero;
9821065
use std::ops::Range;
983-
9841066
#[test]
9851067
fn test_new() {
9861068
Python::with_gil(|py| {
@@ -1523,4 +1605,127 @@ mod tests {
15231605
}
15241606
})
15251607
}
1608+
1609+
#[test]
1610+
fn test_bound_tuple_nth() {
1611+
Python::with_gil(|py| {
1612+
let tuple = PyTuple::new(py, vec![1, 2, 3, 4]).unwrap();
1613+
let mut iter = tuple.iter();
1614+
assert_eq!(iter.nth(1).unwrap().extract::<i32>().unwrap(), 2);
1615+
assert_eq!(iter.nth(1).unwrap().extract::<i32>().unwrap(), 4);
1616+
assert!(iter.nth(1).is_none());
1617+
1618+
let tuple = PyTuple::new(py, Vec::<i32>::new()).unwrap();
1619+
let mut iter = tuple.iter();
1620+
iter.next();
1621+
assert!(iter.nth(1).is_none());
1622+
1623+
let tuple = PyTuple::new(py, vec![1, 2, 3]).unwrap();
1624+
let mut iter = tuple.iter();
1625+
assert!(iter.nth(10).is_none());
1626+
1627+
let tuple = PyTuple::new(py, vec![6, 7, 8, 9, 10]).unwrap();
1628+
let mut iter = tuple.iter();
1629+
assert_eq!(iter.next().unwrap().extract::<i32>().unwrap(), 6);
1630+
assert_eq!(iter.nth(2).unwrap().extract::<i32>().unwrap(), 9);
1631+
assert_eq!(iter.next().unwrap().extract::<i32>().unwrap(), 10);
1632+
1633+
let mut iter = tuple.iter();
1634+
assert_eq!(iter.nth_back(1).unwrap().extract::<i32>().unwrap(), 9);
1635+
assert_eq!(iter.nth(2).unwrap().extract::<i32>().unwrap(), 8);
1636+
assert!(iter.next().is_none());
1637+
});
1638+
}
1639+
1640+
#[test]
1641+
fn test_bound_tuple_nth_back() {
1642+
Python::with_gil(|py| {
1643+
let tuple = PyTuple::new(py, vec![1, 2, 3, 4, 5]).unwrap();
1644+
let mut iter = tuple.iter();
1645+
assert_eq!(iter.nth_back(0).unwrap().extract::<i32>().unwrap(), 5);
1646+
assert_eq!(iter.nth_back(1).unwrap().extract::<i32>().unwrap(), 3);
1647+
assert!(iter.nth_back(2).is_none());
1648+
1649+
let tuple = PyTuple::new(py, Vec::<i32>::new()).unwrap();
1650+
let mut iter = tuple.iter();
1651+
assert!(iter.nth_back(0).is_none());
1652+
assert!(iter.nth_back(1).is_none());
1653+
1654+
let tuple = PyTuple::new(py, vec![1, 2, 3]).unwrap();
1655+
let mut iter = tuple.iter();
1656+
assert!(iter.nth_back(5).is_none());
1657+
1658+
let tuple = PyTuple::new(py, vec![1, 2, 3, 4, 5]).unwrap();
1659+
let mut iter = tuple.iter();
1660+
iter.next_back(); // Consume the last element
1661+
assert_eq!(iter.nth_back(1).unwrap().extract::<i32>().unwrap(), 3);
1662+
assert_eq!(iter.next_back().unwrap().extract::<i32>().unwrap(), 2);
1663+
assert_eq!(iter.nth_back(0).unwrap().extract::<i32>().unwrap(), 1);
1664+
1665+
let tuple = PyTuple::new(py, vec![1, 2, 3, 4, 5]).unwrap();
1666+
let mut iter = tuple.iter();
1667+
assert_eq!(iter.nth_back(1).unwrap().extract::<i32>().unwrap(), 4);
1668+
assert_eq!(iter.nth_back(2).unwrap().extract::<i32>().unwrap(), 1);
1669+
1670+
let mut iter2 = tuple.iter();
1671+
iter2.next_back();
1672+
assert_eq!(iter2.nth_back(1).unwrap().extract::<i32>().unwrap(), 3);
1673+
assert_eq!(iter2.next_back().unwrap().extract::<i32>().unwrap(), 2);
1674+
1675+
let mut iter3 = tuple.iter();
1676+
iter3.nth(1);
1677+
assert_eq!(iter3.nth_back(2).unwrap().extract::<i32>().unwrap(), 3);
1678+
assert!(iter3.nth_back(0).is_none());
1679+
});
1680+
}
1681+
1682+
#[cfg(feature = "nightly")]
1683+
#[test]
1684+
fn test_bound_tuple_advance_by() {
1685+
Python::with_gil(|py| {
1686+
let tuple = PyTuple::new(py, vec![1, 2, 3, 4, 5]).unwrap();
1687+
let mut iter = tuple.iter();
1688+
1689+
assert_eq!(iter.advance_by(2), Ok(()));
1690+
assert_eq!(iter.next().unwrap().extract::<i32>().unwrap(), 3);
1691+
assert_eq!(iter.advance_by(0), Ok(()));
1692+
assert_eq!(iter.advance_by(100), Err(NonZero::new(98).unwrap()));
1693+
assert!(iter.next().is_none());
1694+
1695+
let mut iter2 = tuple.iter();
1696+
assert_eq!(iter2.advance_by(6), Err(NonZero::new(1).unwrap()));
1697+
1698+
let mut iter3 = tuple.iter();
1699+
assert_eq!(iter3.advance_by(5), Ok(()));
1700+
1701+
let mut iter4 = tuple.iter();
1702+
assert_eq!(iter4.advance_by(0), Ok(()));
1703+
assert_eq!(iter4.next().unwrap().extract::<i32>().unwrap(), 1);
1704+
})
1705+
}
1706+
1707+
#[cfg(feature = "nightly")]
1708+
#[test]
1709+
fn test_bound_tuple_advance_back_by() {
1710+
Python::with_gil(|py| {
1711+
let tuple = PyTuple::new(py, vec![1, 2, 3, 4, 5]).unwrap();
1712+
let mut iter = tuple.iter();
1713+
1714+
assert_eq!(iter.advance_back_by(2), Ok(()));
1715+
assert_eq!(iter.next_back().unwrap().extract::<i32>().unwrap(), 3);
1716+
assert_eq!(iter.advance_back_by(0), Ok(()));
1717+
assert_eq!(iter.advance_back_by(100), Err(NonZero::new(98).unwrap()));
1718+
assert!(iter.next_back().is_none());
1719+
1720+
let mut iter2 = tuple.iter();
1721+
assert_eq!(iter2.advance_back_by(6), Err(NonZero::new(1).unwrap()));
1722+
1723+
let mut iter3 = tuple.iter();
1724+
assert_eq!(iter3.advance_back_by(5), Ok(()));
1725+
1726+
let mut iter4 = tuple.iter();
1727+
assert_eq!(iter4.advance_back_by(0), Ok(()));
1728+
assert_eq!(iter4.next_back().unwrap().extract::<i32>().unwrap(), 5);
1729+
})
1730+
}
15261731
}

0 commit comments

Comments
 (0)