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

Optimize nth and nth_back for BoundListIterator #4810

Open
wants to merge 14 commits into
base: main
Choose a base branch
from
Open

Optimize nth and nth_back for BoundListIterator #4810

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
14 commits
Select commit Hold shift + click to select a range
cc9cabd
Optimize nth and nth_back for BoundListIterator. Add unit test and be…
Owen-CH-Leung Dec 21, 2024
3a0c196
Fix fmt and newsfragment CI
Owen-CH-Leung Dec 21, 2024
40d38f3
Fix clippy and changelog CI
Owen-CH-Leung Dec 21, 2024
1b19616
Merge branch 'main' into optimize_nth_and_nthback_for_boundlistiter
Owen-CH-Leung Jan 9, 2025
f6e95a8
Revise Impl of nth and nth_back. Impl advance_by
Owen-CH-Leung Jan 9, 2025
b0c749b
Fix failing fmt
Owen-CH-Leung Jan 9, 2025
b2bf973
Fix failing ruff test
Owen-CH-Leung Jan 9, 2025
4e86709
Merge branch 'main' into optimize_nth_and_nthback_for_boundlistiter
Owen-CH-Leung Jan 10, 2025
e4269c2
branch out nth, nth_unchecked, nth_back, nth_back_unchecked.
Owen-CH-Leung Jan 10, 2025
6e18229
Fix fmt
Owen-CH-Leung Jan 10, 2025
5bab05b
Revise advance_by impl. add advance_by unittest.
Owen-CH-Leung Jan 10, 2025
0b23173
Fix fmt
Owen-CH-Leung Jan 10, 2025
e88f8be
Fix clippy unused function warning
Owen-CH-Leung Jan 10, 2025
3a7a171
Set appropriate Py_LIMITED_API flag
Owen-CH-Leung Jan 10, 2025
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/4810.added.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Optimizes `nth` and `nth_back` for `BoundListIterator`
30 changes: 29 additions & 1 deletion pyo3-benches/benches/bench_list.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,33 @@ fn list_get_item(b: &mut Bencher<'_>) {
});
}

#[cfg(not(any(Py_LIMITED_API, Py_GIL_DISABLED)))]
fn list_nth(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
const LEN: usize = 50;
let list = PyList::new_bound(py, 0..LEN);
let mut sum = 0;
b.iter(|| {
for i in 0..LEN {
sum += list.iter().nth(i).unwrap().extract::<usize>().unwrap();
}
});
});
}

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

#[cfg(not(Py_LIMITED_API))]
fn list_get_item_unchecked(b: &mut Bencher<'_>) {
Python::with_gil(|py| {
const LEN: usize = 50_000;
Expand All @@ -66,6 +92,8 @@ fn sequence_from_list(b: &mut Bencher<'_>) {
fn criterion_benchmark(c: &mut Criterion) {
c.bench_function("iter_list", iter_list);
c.bench_function("list_new", list_new);
c.bench_function("list_nth", list_nth);
c.bench_function("list_nth_back", list_nth_back);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Totally missed this when I added the new free-threaded implementations for folding operations in list. I'll add benchmarks for those in a separate PR.

c.bench_function("list_get_item", list_get_item);
#[cfg(not(any(Py_LIMITED_API, Py_GIL_DISABLED)))]
c.bench_function("list_get_item_unchecked", list_get_item_unchecked);
Expand Down
129 changes: 128 additions & 1 deletion src/types/list.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -494,7 +494,6 @@ impl<'py> Iterator for BoundListIterator<'py> {
#[inline]
fn next(&mut self) -> Option<Self::Item> {
let length = self.length.min(self.list.len());

if self.index < length {
let item = unsafe { self.get_item(self.index) };
self.index += 1;
Expand All @@ -509,6 +508,20 @@ impl<'py> Iterator for BoundListIterator<'py> {
let len = self.len();
(len, Some(len))
}

#[inline]
fn nth(&mut self, n: usize) -> Option<Self::Item> {
let length = self.length.min(self.list.len());
let target_index = self.index + n;
if self.index + n < length {
let item = unsafe { self.get_item(target_index) };
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder, is there a time-of-check to time-of-use bug here on the length? Not one for this PR, but a follow up I will try not to forget...

Copy link
Contributor Author

@Owen-CH-Leung Owen-CH-Leung Dec 29, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yup I think that's a potential TOCTOU bug here. This particular implementation assumes the user ensures proper synchronization if they intend to use the iterator in a multi-threaded or mutable environment.

Copy link
Contributor Author

@Owen-CH-Leung Owen-CH-Leung Dec 29, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think the current implementation of next also has a potential TOCTOU bug:

https://github.com/PyO3/pyo3/blob/main/src/types/list.rs#L495-L498

self.index = target_index + 1;
Some(item)
} else {
self.index = self.list.len();
None
}
}
}

impl DoubleEndedIterator for BoundListIterator<'_> {
Expand All @@ -524,6 +537,20 @@ impl DoubleEndedIterator for BoundListIterator<'_> {
None
}
}

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

impl ExactSizeIterator for BoundListIterator<'_> {
Expand Down Expand Up @@ -720,6 +747,106 @@ mod tests {
});
}

#[test]
fn test_iter_nth() {
Python::with_gil(|py| {
let v = vec![6, 7, 8, 9, 10];
let ob = (&v).into_pyobject(py).unwrap();
let list = ob.downcast::<PyList>().unwrap();

let mut iter = list.iter();
iter.next();
assert_eq!(iter.nth(1).unwrap().extract::<i32>().unwrap(), 8);
assert_eq!(iter.nth(1).unwrap().extract::<i32>().unwrap(), 10);
assert!(iter.nth(1).is_none());

let v: Vec<i32> = vec![];
let ob = (&v).into_pyobject(py).unwrap();
let list = ob.downcast::<PyList>().unwrap();

let mut iter = list.iter();
iter.next();
assert!(iter.nth(1).is_none());

let v = vec![1, 2, 3];
let ob = (&v).into_pyobject(py).unwrap();
let list = ob.downcast::<PyList>().unwrap();

let mut iter = list.iter();
assert!(iter.nth(10).is_none());

let v = vec![6, 7, 8, 9, 10];
let ob = (&v).into_pyobject(py).unwrap();
let list = ob.downcast::<PyList>().unwrap();
let mut iter = list.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 = list.iter();
iter.nth_back(1);
assert_eq!(iter.nth(2).unwrap().extract::<i32>().unwrap(), 8);
assert!(iter.next().is_none());
});
}

#[test]
fn test_iter_nth_back() {
Python::with_gil(|py| {
let v = vec![1, 2, 3, 4, 5];
let ob = (&v).into_pyobject(py).unwrap();
let list = ob.downcast::<PyList>().unwrap();

let mut iter = list.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 v: Vec<i32> = vec![];
let ob = (&v).into_pyobject(py).unwrap();
let list = ob.downcast::<PyList>().unwrap();

let mut iter = list.iter();
assert!(iter.nth_back(0).is_none());
assert!(iter.nth_back(1).is_none());

let v = vec![1, 2, 3];
let ob = (&v).into_pyobject(py).unwrap();
let list = ob.downcast::<PyList>().unwrap();

let mut iter = list.iter();
assert!(iter.nth_back(5).is_none());

let v = vec![1, 2, 3, 4, 5];
let ob = (&v).into_pyobject(py).unwrap();
let list = ob.downcast::<PyList>().unwrap();

let mut iter = list.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 v = vec![1, 2, 3, 4, 5];
let ob = (&v).into_pyobject(py).unwrap();
let list = ob.downcast::<PyList>().unwrap();

let mut iter = list.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 = list.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 = list.iter();
iter3.nth(1);
assert_eq!(iter3.nth_back(2).unwrap().extract::<i32>().unwrap(), 3);
assert!(iter3.nth_back(0).is_none());
});
}

#[test]
fn test_iter_rev() {
Python::with_gil(|py| {
Expand Down
Loading