From cc9cabd30fe76d4340941fa3bbc23ec39eda70d4 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Sat, 21 Dec 2024 14:43:58 +0800 Subject: [PATCH 01/19] Optimize nth and nth_back for BoundListIterator. Add unit test and benchmarks --- pyo3-benches/benches/bench_list.rs | 30 ++++++- src/types/list.rs | 137 ++++++++++++++++++++++++++++- 2 files changed, 165 insertions(+), 2 deletions(-) diff --git a/pyo3-benches/benches/bench_list.rs b/pyo3-benches/benches/bench_list.rs index cc790db37bf..7a19452455e 100644 --- a/pyo3-benches/benches/bench_list.rs +++ b/pyo3-benches/benches/bench_list.rs @@ -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::().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::().unwrap(); + } + }); + }); +} + +#[cfg(not(Py_LIMITED_API))] fn list_get_item_unchecked(b: &mut Bencher<'_>) { Python::with_gil(|py| { const LEN: usize = 50_000; @@ -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); 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); diff --git a/src/types/list.rs b/src/types/list.rs index af2b557cba9..07089872bae 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -494,7 +494,6 @@ impl<'py> Iterator for BoundListIterator<'py> { #[inline] fn next(&mut self) -> Option { let length = self.length.min(self.list.len()); - if self.index < length { let item = unsafe { self.get_item(self.index) }; self.index += 1; @@ -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 { + 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) }; + self.index = target_index + 1; + Some(item) + } else { + self.index = self.list.len(); + None + } + } } impl DoubleEndedIterator for BoundListIterator<'_> { @@ -524,6 +537,20 @@ impl DoubleEndedIterator for BoundListIterator<'_> { None } } + + #[inline] + fn nth_back(&mut self, n: usize) -> Option { + 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<'_> { @@ -720,6 +747,114 @@ 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::().unwrap(); + + let mut iter = list.iter(); + assert_eq!(iter.nth(0).unwrap().extract::().unwrap(), 6); + assert_eq!(iter.nth(1).unwrap().extract::().unwrap(), 8); + assert_eq!(iter.nth(1).unwrap().extract::().unwrap(), 10); + assert!(iter.nth(1).is_none()); + + let v: Vec = vec![]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().unwrap(); + + let mut iter = list.iter(); + assert!(iter.nth(0).is_none()); + assert!(iter.nth(1).is_none()); + + let v = vec![1, 2, 3]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().unwrap(); + + let mut iter = list.iter(); + assert!(iter.nth(10).is_none()); + + let v = vec![10]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().unwrap(); + + let mut iter = list.iter(); + assert_eq!(iter.nth(0).unwrap().extract::().unwrap(), 10); + assert!(iter.nth(0).is_none()); + + let v = vec![6, 7, 8, 9, 10]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().unwrap(); + let mut iter = list.iter(); + assert_eq!(iter.next().unwrap().extract::().unwrap(), 6); + assert_eq!(iter.nth(2).unwrap().extract::().unwrap(), 9); + assert_eq!(iter.next().unwrap().extract::().unwrap(), 10); + + let mut iter = list.iter(); + iter.nth_back(1); + assert_eq!(iter.nth(2).unwrap().extract::().unwrap(), 8); + assert!(iter.nth(0).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::().unwrap(); + + let mut iter = list.iter(); + assert_eq!(iter.nth_back(0).unwrap().extract::().unwrap(), 5); + assert_eq!(iter.nth_back(1).unwrap().extract::().unwrap(), 3); + assert!(iter.nth_back(2).is_none()); + + let v: Vec = vec![]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().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::().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::().unwrap(); + + let mut iter = list.iter(); + iter.next_back(); // Consume the last element + assert_eq!(iter.nth_back(1).unwrap().extract::().unwrap(), 3); + assert_eq!(iter.next_back().unwrap().extract::().unwrap(), 2); + assert_eq!(iter.nth_back(0).unwrap().extract::().unwrap(), 1); + + let v = vec![1,2,3,4,5]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().unwrap(); + + let mut iter = list.iter(); + assert_eq!(iter.nth_back(1).unwrap().extract::().unwrap(), 4); + assert_eq!(iter.nth_back(2).unwrap().extract::().unwrap(), 1); + + let mut iter2 = list.iter(); + iter2.next_back(); + assert_eq!(iter2.nth_back(1).unwrap().extract::().unwrap(), 3); + assert_eq!(iter2.next_back().unwrap().extract::().unwrap(), 2); + + let mut iter3 = list.iter(); + iter3.nth(1); + assert_eq!(iter3.nth_back(2).unwrap().extract::().unwrap(), 3); + assert!(iter3.nth_back(0).is_none()); + }); + } + #[test] fn test_iter_rev() { Python::with_gil(|py| { From 3a0c19665c368e190ada74376d89ad86ba087dc7 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Sat, 21 Dec 2024 16:09:23 +0800 Subject: [PATCH 02/19] Fix fmt and newsfragment CI --- newsfragments/4787.added.md | 1 + src/types/list.rs | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) create mode 100644 newsfragments/4787.added.md diff --git a/newsfragments/4787.added.md b/newsfragments/4787.added.md new file mode 100644 index 00000000000..e89e22e544d --- /dev/null +++ b/newsfragments/4787.added.md @@ -0,0 +1 @@ +Optimizes `nth` and `nth_back` for `BoundListIterator` \ No newline at end of file diff --git a/src/types/list.rs b/src/types/list.rs index 07089872bae..9e74d7ad1d6 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -835,7 +835,7 @@ mod tests { assert_eq!(iter.next_back().unwrap().extract::().unwrap(), 2); assert_eq!(iter.nth_back(0).unwrap().extract::().unwrap(), 1); - let v = vec![1,2,3,4,5]; + let v = vec![1, 2, 3, 4, 5]; let ob = (&v).into_pyobject(py).unwrap(); let list = ob.downcast::().unwrap(); From 40d38f3c66fa0488f8d814c22a303354eb24a2c9 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Sat, 21 Dec 2024 16:55:29 +0800 Subject: [PATCH 03/19] Fix clippy and changelog CI --- newsfragments/{4787.added.md => 4810.added.md} | 0 src/types/list.rs | 14 +++----------- 2 files changed, 3 insertions(+), 11 deletions(-) rename newsfragments/{4787.added.md => 4810.added.md} (100%) diff --git a/newsfragments/4787.added.md b/newsfragments/4810.added.md similarity index 100% rename from newsfragments/4787.added.md rename to newsfragments/4810.added.md diff --git a/src/types/list.rs b/src/types/list.rs index 9e74d7ad1d6..3db72d07fed 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -755,7 +755,7 @@ mod tests { let list = ob.downcast::().unwrap(); let mut iter = list.iter(); - assert_eq!(iter.nth(0).unwrap().extract::().unwrap(), 6); + iter.next(); assert_eq!(iter.nth(1).unwrap().extract::().unwrap(), 8); assert_eq!(iter.nth(1).unwrap().extract::().unwrap(), 10); assert!(iter.nth(1).is_none()); @@ -765,7 +765,7 @@ mod tests { let list = ob.downcast::().unwrap(); let mut iter = list.iter(); - assert!(iter.nth(0).is_none()); + iter.next(); assert!(iter.nth(1).is_none()); let v = vec![1, 2, 3]; @@ -775,14 +775,6 @@ mod tests { let mut iter = list.iter(); assert!(iter.nth(10).is_none()); - let v = vec![10]; - let ob = (&v).into_pyobject(py).unwrap(); - let list = ob.downcast::().unwrap(); - - let mut iter = list.iter(); - assert_eq!(iter.nth(0).unwrap().extract::().unwrap(), 10); - assert!(iter.nth(0).is_none()); - let v = vec![6, 7, 8, 9, 10]; let ob = (&v).into_pyobject(py).unwrap(); let list = ob.downcast::().unwrap(); @@ -794,7 +786,7 @@ mod tests { let mut iter = list.iter(); iter.nth_back(1); assert_eq!(iter.nth(2).unwrap().extract::().unwrap(), 8); - assert!(iter.nth(0).is_none()); + assert!(iter.next().is_none()); }); } From f6e95a8d9b127af08104f9b7ea0ebf8e07b5ecd9 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Thu, 9 Jan 2025 23:12:31 +0800 Subject: [PATCH 04/19] Revise Impl of nth and nth_back. Impl advance_by --- src/types/list.rs | 188 +++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 187 insertions(+), 1 deletion(-) diff --git a/src/types/list.rs b/src/types/list.rs index 76da36d00b9..21bfb4d9960 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -1,5 +1,4 @@ use std::iter::FusedIterator; - use crate::err::{self, PyResult}; use crate::ffi::{self, Py_ssize_t}; use crate::ffi_ptr_ext::FfiPtrExt; @@ -547,6 +546,31 @@ impl<'py> BoundListIterator<'py> { } } + /// # Safety + /// + /// On the free-threaded build, caller must verify they have exclusive + /// access to the list by holding a lock or by holding the innermost + /// critical section on the list. + #[inline] + #[cfg(not(Py_LIMITED_API))] + #[deny(unsafe_op_in_unsafe_fn)] + unsafe fn nth_unchecked( + index: &mut Index, + length: &mut Length, + list: &Bound<'py, PyList>, + n: usize, + ) -> Option> { + let length = length.0.min(list.len()); + let target_index = index.0 + n; + if index.0 + n < length { + let item = unsafe { list.get_item_unchecked(target_index) }; + index.0 = target_index + 1; + Some(item) + } else { + None + } + } + /// # Safety /// /// On the free-threaded build, caller must verify they have exclusive @@ -589,6 +613,31 @@ impl<'py> BoundListIterator<'py> { } } + /// # Safety + /// + /// On the free-threaded build, caller must verify they have exclusive + /// access to the list by holding a lock or by holding the innermost + /// critical section on the list. + #[inline] + #[cfg(not(Py_LIMITED_API))] + #[deny(unsafe_op_in_unsafe_fn)] + unsafe fn nth_back_unchecked( + index: &mut Index, + length: &mut Length, + list: &Bound<'py, PyList>, + n: usize, + ) -> Option> { + let length_size = length.0.min(list.len()); + if index.0 + n < length_size { + let target_index = length_size - n - 1; + let item = unsafe {list.get_item_unchecked(target_index)}; + *length = Length(target_index); + Some(item) + } else { + None + } + } + #[cfg(not(Py_LIMITED_API))] fn with_critical_section( &mut self, @@ -625,6 +674,14 @@ impl<'py> Iterator for BoundListIterator<'py> { } } + #[inline] + #[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))] + fn nth(&mut self, n: usize) -> Option { + self.with_critical_section(|index, length, list| unsafe { + Self::nth_unchecked(index, length, list, n) + }) + } + #[inline] fn size_hint(&self) -> (usize, Option) { let len = self.len(); @@ -750,6 +807,27 @@ impl<'py> Iterator for BoundListIterator<'py> { None }) } + + #[inline] + #[cfg(all(Py_GIL_DISABLED, feature = "nightly"))] + fn advance_by(&mut self, n: usize) -> Result<(), NonZero> { + self.with_critical_section(|index, length, list| { + let length = length.0.min(list.len()); + let target_index = index.0 + n; + if index.0 + n < length { + let item = list.get_item(target_index); + match item { + Ok(_) => { + index.0 = target_index; + Ok(()) + } + Err(_) => Err(NonZero::new(n - index.0)) + } + } else { + Err(NonZero::new(n - index.0)) + } + }) + } } impl DoubleEndedIterator for BoundListIterator<'_> { @@ -772,6 +850,14 @@ impl DoubleEndedIterator for BoundListIterator<'_> { } } + #[inline] + #[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))] + fn nth_back(&mut self, n: usize) -> Option { + self.with_critical_section(|index, length, list| unsafe { + Self::nth_back_unchecked(index, length, list, n) + }) + } + #[inline] #[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))] fn rfold(mut self, init: B, mut f: F) -> B @@ -1502,4 +1588,104 @@ mod tests { assert!(tuple.eq(tuple_expected).unwrap()); }) } + + #[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::().unwrap(); + + let mut iter = list.iter(); + iter.next(); + assert_eq!(iter.nth(1).unwrap().extract::().unwrap(), 8); + assert_eq!(iter.nth(1).unwrap().extract::().unwrap(), 10); + assert!(iter.nth(1).is_none()); + + let v: Vec = vec![]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().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::().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::().unwrap(); + let mut iter = list.iter(); + assert_eq!(iter.next().unwrap().extract::().unwrap(), 6); + assert_eq!(iter.nth(2).unwrap().extract::().unwrap(), 9); + assert_eq!(iter.next().unwrap().extract::().unwrap(), 10); + + let mut iter = list.iter(); + iter.nth_back(1); + assert_eq!(iter.nth(2).unwrap().extract::().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::().unwrap(); + + let mut iter = list.iter(); + assert_eq!(iter.nth_back(0).unwrap().extract::().unwrap(), 5); + assert_eq!(iter.nth_back(1).unwrap().extract::().unwrap(), 3); + assert!(iter.nth_back(2).is_none()); + + let v: Vec = vec![]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().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::().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::().unwrap(); + + let mut iter = list.iter(); + iter.next_back(); // Consume the last element + assert_eq!(iter.nth_back(1).unwrap().extract::().unwrap(), 3); + assert_eq!(iter.next_back().unwrap().extract::().unwrap(), 2); + assert_eq!(iter.nth_back(0).unwrap().extract::().unwrap(), 1); + + let v = vec![1, 2, 3, 4, 5]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().unwrap(); + + let mut iter = list.iter(); + assert_eq!(iter.nth_back(1).unwrap().extract::().unwrap(), 4); + assert_eq!(iter.nth_back(2).unwrap().extract::().unwrap(), 1); + + let mut iter2 = list.iter(); + iter2.next_back(); + assert_eq!(iter2.nth_back(1).unwrap().extract::().unwrap(), 3); + assert_eq!(iter2.next_back().unwrap().extract::().unwrap(), 2); + + let mut iter3 = list.iter(); + iter3.nth(1); + assert_eq!(iter3.nth_back(2).unwrap().extract::().unwrap(), 3); + assert!(iter3.nth_back(0).is_none()); + }); + } } From b0c749bf56a770a8dc7056bc9d80587460a978a8 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Thu, 9 Jan 2025 23:14:26 +0800 Subject: [PATCH 05/19] Fix failing fmt --- src/types/list.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index 21bfb4d9960..ab58ab5f033 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -1,4 +1,3 @@ -use std::iter::FusedIterator; use crate::err::{self, PyResult}; use crate::ffi::{self, Py_ssize_t}; use crate::ffi_ptr_ext::FfiPtrExt; @@ -7,6 +6,7 @@ use crate::types::{PySequence, PyTuple}; use crate::{ Borrowed, Bound, BoundObject, IntoPyObject, IntoPyObjectExt, PyAny, PyErr, PyObject, Python, }; +use std::iter::FusedIterator; use crate::types::any::PyAnyMethods; use crate::types::sequence::PySequenceMethods; @@ -630,7 +630,7 @@ impl<'py> BoundListIterator<'py> { let length_size = length.0.min(list.len()); if index.0 + n < length_size { let target_index = length_size - n - 1; - let item = unsafe {list.get_item_unchecked(target_index)}; + let item = unsafe { list.get_item_unchecked(target_index) }; *length = Length(target_index); Some(item) } else { @@ -821,7 +821,7 @@ impl<'py> Iterator for BoundListIterator<'py> { index.0 = target_index; Ok(()) } - Err(_) => Err(NonZero::new(n - index.0)) + Err(_) => Err(NonZero::new(n - index.0)), } } else { Err(NonZero::new(n - index.0)) From b2bf973935fc110bf1013356a7fb631ba284896d Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Thu, 9 Jan 2025 23:23:54 +0800 Subject: [PATCH 06/19] Fix failing ruff test --- noxfile.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/noxfile.py b/noxfile.py index 25cb8d1eb92..760c21d03df 100644 --- a/noxfile.py +++ b/noxfile.py @@ -351,7 +351,7 @@ def test_emscripten(session: nox.Session): session, "bash", "-c", - f"source {info.builddir/'emsdk/emsdk_env.sh'} && cargo test", + f"source {info.builddir / 'emsdk/emsdk_env.sh'} && cargo test", ) @@ -951,7 +951,7 @@ def set( f"""\ implementation={implementation} version={version} -build_flags={','.join(build_flags)} +build_flags={",".join(build_flags)} suppress_build_script_link_lines=true """ ) From e4269c273453562ed65178cd47c341042119409c Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Fri, 10 Jan 2025 10:19:20 +0800 Subject: [PATCH 07/19] branch out nth, nth_unchecked, nth_back, nth_back_unchecked. --- src/types/list.rs | 101 +++++++++++++++++++++++++++++++++------------- 1 file changed, 72 insertions(+), 29 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index ab58ab5f033..a99a0e6a17a 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -7,7 +7,8 @@ use crate::{ Borrowed, Bound, BoundObject, IntoPyObject, IntoPyObjectExt, PyAny, PyErr, PyObject, Python, }; use std::iter::FusedIterator; - +#[cfg(feature = "nightly")] +use std::num::NonZero; use crate::types::any::PyAnyMethods; use crate::types::sequence::PySequenceMethods; @@ -546,11 +547,6 @@ impl<'py> BoundListIterator<'py> { } } - /// # Safety - /// - /// On the free-threaded build, caller must verify they have exclusive - /// access to the list by holding a lock or by holding the innermost - /// critical section on the list. #[inline] #[cfg(not(Py_LIMITED_API))] #[deny(unsafe_op_in_unsafe_fn)] @@ -571,6 +567,26 @@ impl<'py> BoundListIterator<'py> { } } + #[inline] + #[cfg(Py_LIMITED_API)] + #[deny(unsafe_op_in_unsafe_fn)] + fn nth( + index: &mut Index, + length: &mut Length, + list: &Bound<'py, PyList>, + n: usize, + ) -> Option> { + let length = length.0.min(list.len()); + let target_index = index.0 + n; + if index.0 + n < length { + let item =list.get_item(target_index).expect("get-item failed"); + index.0 = target_index + 1; + Some(item) + } else { + None + } + } + /// # Safety /// /// On the free-threaded build, caller must verify they have exclusive @@ -613,11 +629,6 @@ impl<'py> BoundListIterator<'py> { } } - /// # Safety - /// - /// On the free-threaded build, caller must verify they have exclusive - /// access to the list by holding a lock or by holding the innermost - /// critical section on the list. #[inline] #[cfg(not(Py_LIMITED_API))] #[deny(unsafe_op_in_unsafe_fn)] @@ -638,6 +649,25 @@ impl<'py> BoundListIterator<'py> { } } + #[inline] + #[cfg(Py_LIMITED_API)] + fn nth_back( + index: &mut Index, + length: &mut Length, + list: &Bound<'py, PyList>, + n: usize, + ) -> Option> { + let length_size = length.0.min(list.len()); + if index.0 + n < length_size { + let target_index = length_size - n - 1; + let item = list.get_item(target_index).expect("get-item failed"); + *length = Length(target_index); + Some(item) + } else { + None + } + } + #[cfg(not(Py_LIMITED_API))] fn with_critical_section( &mut self, @@ -675,11 +705,20 @@ impl<'py> Iterator for BoundListIterator<'py> { } #[inline] - #[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))] fn nth(&mut self, n: usize) -> Option { - self.with_critical_section(|index, length, list| unsafe { - Self::nth_unchecked(index, length, list, n) - }) + #[cfg(not(Py_LIMITED_API))] { + self.with_critical_section(|index, length, list| unsafe { + Self::nth_unchecked(index, length, list, n) + }) + } + #[cfg(Py_LIMITED_API)] { + let Self { + index, + length, + list, + } = self; + Self::nth(index, length, list, n) + } } #[inline] @@ -809,22 +848,17 @@ impl<'py> Iterator for BoundListIterator<'py> { } #[inline] - #[cfg(all(Py_GIL_DISABLED, feature = "nightly"))] + #[cfg(feature = "nightly")] fn advance_by(&mut self, n: usize) -> Result<(), NonZero> { self.with_critical_section(|index, length, list| { let length = length.0.min(list.len()); let target_index = index.0 + n; if index.0 + n < length { - let item = list.get_item(target_index); - match item { - Ok(_) => { - index.0 = target_index; - Ok(()) - } - Err(_) => Err(NonZero::new(n - index.0)), - } + let item = unsafe { list.get_item_unchecked(target_index) }; + index.0 = target_index; + Ok(()) } else { - Err(NonZero::new(n - index.0)) + Err(unsafe { NonZero::new_unchecked(n - index.0) }) } }) } @@ -851,11 +885,20 @@ impl DoubleEndedIterator for BoundListIterator<'_> { } #[inline] - #[cfg(all(Py_GIL_DISABLED, not(feature = "nightly")))] fn nth_back(&mut self, n: usize) -> Option { - self.with_critical_section(|index, length, list| unsafe { - Self::nth_back_unchecked(index, length, list, n) - }) + #[cfg(not(Py_LIMITED_API))] { + self.with_critical_section(|index, length, list| unsafe { + Self::nth_back_unchecked(index, length, list, n) + }) + } + #[cfg(Py_LIMITED_API)] { + let Self { + index, + length, + list, + } = self; + Self::nth_back(index, length, list, n) + } } #[inline] From 6e182291715bd55b2f67ef4b6ccb60339009f26b Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Fri, 10 Jan 2025 10:22:59 +0800 Subject: [PATCH 08/19] Fix fmt --- src/types/list.rs | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index a99a0e6a17a..56dab96dba2 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -2,6 +2,8 @@ use crate::err::{self, PyResult}; use crate::ffi::{self, Py_ssize_t}; use crate::ffi_ptr_ext::FfiPtrExt; use crate::internal_tricks::get_ssize_index; +use crate::types::any::PyAnyMethods; +use crate::types::sequence::PySequenceMethods; use crate::types::{PySequence, PyTuple}; use crate::{ Borrowed, Bound, BoundObject, IntoPyObject, IntoPyObjectExt, PyAny, PyErr, PyObject, Python, @@ -9,8 +11,6 @@ use crate::{ use std::iter::FusedIterator; #[cfg(feature = "nightly")] use std::num::NonZero; -use crate::types::any::PyAnyMethods; -use crate::types::sequence::PySequenceMethods; /// Represents a Python `list`. /// @@ -579,7 +579,7 @@ impl<'py> BoundListIterator<'py> { let length = length.0.min(list.len()); let target_index = index.0 + n; if index.0 + n < length { - let item =list.get_item(target_index).expect("get-item failed"); + let item = list.get_item(target_index).expect("get-item failed"); index.0 = target_index + 1; Some(item) } else { @@ -706,12 +706,14 @@ impl<'py> Iterator for BoundListIterator<'py> { #[inline] fn nth(&mut self, n: usize) -> Option { - #[cfg(not(Py_LIMITED_API))] { + #[cfg(not(Py_LIMITED_API))] + { self.with_critical_section(|index, length, list| unsafe { Self::nth_unchecked(index, length, list, n) }) } - #[cfg(Py_LIMITED_API)] { + #[cfg(Py_LIMITED_API)] + { let Self { index, length, @@ -886,12 +888,14 @@ impl DoubleEndedIterator for BoundListIterator<'_> { #[inline] fn nth_back(&mut self, n: usize) -> Option { - #[cfg(not(Py_LIMITED_API))] { + #[cfg(not(Py_LIMITED_API))] + { self.with_critical_section(|index, length, list| unsafe { Self::nth_back_unchecked(index, length, list, n) }) } - #[cfg(Py_LIMITED_API)] { + #[cfg(Py_LIMITED_API)] + { let Self { index, length, From 5bab05b94e04643a2072d00a8699cd307f023cc8 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Fri, 10 Jan 2025 18:53:54 +0800 Subject: [PATCH 09/19] Revise advance_by impl. add advance_by unittest. --- src/lib.rs | 2 +- src/types/list.rs | 53 +++++++++++++++++++++++++++++++++++++++++------ 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index e5146c81c00..3547c52e57c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,7 @@ #![warn(missing_docs)] #![cfg_attr( feature = "nightly", - feature(auto_traits, negative_impls, try_trait_v2) + feature(auto_traits, negative_impls, try_trait_v2, iter_advance_by) )] #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))] // Deny some lints in doctests. diff --git a/src/types/list.rs b/src/types/list.rs index 56dab96dba2..836de2ad5cf 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -705,6 +705,7 @@ impl<'py> Iterator for BoundListIterator<'py> { } #[inline] + #[cfg(not(feature = "nightly"))] fn nth(&mut self, n: usize) -> Option { #[cfg(not(Py_LIMITED_API))] { @@ -853,14 +854,24 @@ impl<'py> Iterator for BoundListIterator<'py> { #[cfg(feature = "nightly")] fn advance_by(&mut self, n: usize) -> Result<(), NonZero> { self.with_critical_section(|index, length, list| { - let length = length.0.min(list.len()); - let target_index = index.0 + n; - if index.0 + n < length { - let item = unsafe { list.get_item_unchecked(target_index) }; - index.0 = target_index; + let max_len = length.0.min(list.len()); + let currently_at = index.0; + 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 { + index.0 += n; Ok(()) } else { - Err(unsafe { NonZero::new_unchecked(n - index.0) }) + index.0 = max_len; + let remainder = n - items_left; + Err(unsafe { NonZero::new_unchecked(remainder) }) } }) } @@ -887,6 +898,7 @@ impl DoubleEndedIterator for BoundListIterator<'_> { } #[inline] + #[cfg(not(feature = "nightly"))] fn nth_back(&mut self, n: usize) -> Option { #[cfg(not(Py_LIMITED_API))] { @@ -967,6 +979,8 @@ impl<'py> IntoIterator for &Bound<'py, PyList> { #[cfg(test)] mod tests { + #[cfg(feature = "nightly")] + use std::num::NonZero; use crate::types::any::PyAnyMethods; use crate::types::list::PyListMethods; use crate::types::sequence::PySequenceMethods; @@ -1735,4 +1749,31 @@ mod tests { assert!(iter3.nth_back(0).is_none()); }); } + + #[cfg(feature = "nightly")] + #[test] + fn test_iter_advance_by() { + Python::with_gil(|py| { + let v = vec![1, 2, 3, 4, 5]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().unwrap(); + + let mut iter = list.iter(); + assert_eq!(iter.advance_by(2), Ok(())); + assert_eq!(iter.next().unwrap().extract::().unwrap(), 3); + assert_eq!(iter.advance_by(0), Ok(())); + assert_eq!(iter.advance_by(100), Err(NonZero::new(98).unwrap())); + + let mut iter2 = list.iter(); + assert_eq!(iter2.advance_by(6), Err(NonZero::new(1).unwrap())); + + let mut iter3 = list.iter(); + assert_eq!(iter3.advance_by(5), Ok(())); + + let mut iter4 = list.iter(); + assert_eq!(iter4.advance_by(0), Ok(())); + assert_eq!(iter4.next().unwrap().extract::().unwrap(), 1); + }) + } + } From 0b23173a62049d8fbadcc63d4850b7eac40a318f Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Fri, 10 Jan 2025 18:57:05 +0800 Subject: [PATCH 10/19] Fix fmt --- src/types/list.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index 836de2ad5cf..5d06205db7f 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -979,13 +979,13 @@ impl<'py> IntoIterator for &Bound<'py, PyList> { #[cfg(test)] mod tests { - #[cfg(feature = "nightly")] - use std::num::NonZero; use crate::types::any::PyAnyMethods; use crate::types::list::PyListMethods; use crate::types::sequence::PySequenceMethods; use crate::types::{PyList, PyTuple}; use crate::{ffi, IntoPyObject, PyResult, Python}; + #[cfg(feature = "nightly")] + use std::num::NonZero; #[test] fn test_new() { @@ -1775,5 +1775,4 @@ mod tests { assert_eq!(iter4.next().unwrap().extract::().unwrap(), 1); }) } - } From e88f8bee7c2f232391184a30e6a3d4be7567188c Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Fri, 10 Jan 2025 20:16:11 +0800 Subject: [PATCH 11/19] Fix clippy unused function warning --- src/types/list.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index 5d06205db7f..2ca59813c24 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -548,7 +548,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(not(Py_LIMITED_API))] + #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] #[deny(unsafe_op_in_unsafe_fn)] unsafe fn nth_unchecked( index: &mut Index, @@ -630,7 +630,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(not(Py_LIMITED_API))] + #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] #[deny(unsafe_op_in_unsafe_fn)] unsafe fn nth_back_unchecked( index: &mut Index, From 3a7a171e0f149ff10438fde43c7f8d0778dbd2e3 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Fri, 10 Jan 2025 21:06:53 +0800 Subject: [PATCH 12/19] Set appropriate Py_LIMITED_API flag --- src/types/list.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index 2ca59813c24..8391696f003 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -9,7 +9,7 @@ use crate::{ Borrowed, Bound, BoundObject, IntoPyObject, IntoPyObjectExt, PyAny, PyErr, PyObject, Python, }; use std::iter::FusedIterator; -#[cfg(feature = "nightly")] +#[cfg(all(not(Py_LIMITED_API), feature = "nightly"))] use std::num::NonZero; /// Represents a Python `list`. @@ -548,7 +548,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] + #[cfg(all(not(Py_LIMITED_API), feature = "nightly"))] #[deny(unsafe_op_in_unsafe_fn)] unsafe fn nth_unchecked( index: &mut Index, @@ -568,7 +568,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(Py_LIMITED_API)] + #[cfg(all(Py_LIMITED_API, feature = "nightly"))] #[deny(unsafe_op_in_unsafe_fn)] fn nth( index: &mut Index, @@ -630,7 +630,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] + #[cfg(all(not(Py_LIMITED_API), feature = "nightly"))] #[deny(unsafe_op_in_unsafe_fn)] unsafe fn nth_back_unchecked( index: &mut Index, @@ -650,7 +650,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(Py_LIMITED_API)] + #[cfg(all(Py_LIMITED_API, feature = "nightly"))] fn nth_back( index: &mut Index, length: &mut Length, @@ -705,7 +705,7 @@ impl<'py> Iterator for BoundListIterator<'py> { } #[inline] - #[cfg(not(feature = "nightly"))] + #[cfg(feature = "nightly")] fn nth(&mut self, n: usize) -> Option { #[cfg(not(Py_LIMITED_API))] { @@ -851,7 +851,7 @@ impl<'py> Iterator for BoundListIterator<'py> { } #[inline] - #[cfg(feature = "nightly")] + #[cfg(all(not(Py_LIMITED_API), feature = "nightly"))] fn advance_by(&mut self, n: usize) -> Result<(), NonZero> { self.with_critical_section(|index, length, list| { let max_len = length.0.min(list.len()); @@ -898,7 +898,7 @@ impl DoubleEndedIterator for BoundListIterator<'_> { } #[inline] - #[cfg(not(feature = "nightly"))] + #[cfg(feature = "nightly")] fn nth_back(&mut self, n: usize) -> Option { #[cfg(not(Py_LIMITED_API))] { From ed8dba63682b3f773ee623a1f7c2320934a12413 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Tue, 14 Jan 2025 22:12:25 +0800 Subject: [PATCH 13/19] Rewrite nth & nth_back using conditional compilation. Rearrange flags for proper compilation --- src/types/list.rs | 109 +++++++++++++--------------------------------- 1 file changed, 30 insertions(+), 79 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index 8391696f003..46ac31b5254 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -9,7 +9,7 @@ use crate::{ Borrowed, Bound, BoundObject, IntoPyObject, IntoPyObjectExt, PyAny, PyErr, PyObject, Python, }; use std::iter::FusedIterator; -#[cfg(all(not(Py_LIMITED_API), feature = "nightly"))] +#[cfg(feature = "nightly")] use std::num::NonZero; /// Represents a Python `list`. @@ -548,28 +548,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(all(not(Py_LIMITED_API), feature = "nightly"))] - #[deny(unsafe_op_in_unsafe_fn)] - unsafe fn nth_unchecked( - index: &mut Index, - length: &mut Length, - list: &Bound<'py, PyList>, - n: usize, - ) -> Option> { - let length = length.0.min(list.len()); - let target_index = index.0 + n; - if index.0 + n < length { - let item = unsafe { list.get_item_unchecked(target_index) }; - index.0 = target_index + 1; - Some(item) - } else { - None - } - } - - #[inline] - #[cfg(all(Py_LIMITED_API, feature = "nightly"))] - #[deny(unsafe_op_in_unsafe_fn)] + #[cfg(not(feature = "nightly"))] fn nth( index: &mut Index, length: &mut Length, @@ -578,8 +557,16 @@ impl<'py> BoundListIterator<'py> { ) -> Option> { let length = length.0.min(list.len()); let target_index = index.0 + n; - if index.0 + n < length { - let item = list.get_item(target_index).expect("get-item failed"); + if target_index < length { + let item = { + #[cfg(Py_LIMITED_API)] { + list.get_item(target_index).expect("get-item failed") + } + + #[cfg(not(Py_LIMITED_API))] { + unsafe { list.get_item_unchecked(target_index) } + } + }; index.0 = target_index + 1; Some(item) } else { @@ -630,27 +617,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(all(not(Py_LIMITED_API), feature = "nightly"))] - #[deny(unsafe_op_in_unsafe_fn)] - unsafe fn nth_back_unchecked( - index: &mut Index, - length: &mut Length, - list: &Bound<'py, PyList>, - n: usize, - ) -> Option> { - let length_size = length.0.min(list.len()); - if index.0 + n < length_size { - let target_index = length_size - n - 1; - let item = unsafe { list.get_item_unchecked(target_index) }; - *length = Length(target_index); - Some(item) - } else { - None - } - } - - #[inline] - #[cfg(all(Py_LIMITED_API, feature = "nightly"))] + #[cfg(not(feature = "nightly"))] fn nth_back( index: &mut Index, length: &mut Length, @@ -660,7 +627,17 @@ impl<'py> BoundListIterator<'py> { let length_size = length.0.min(list.len()); if index.0 + n < length_size { let target_index = length_size - n - 1; - let item = list.get_item(target_index).expect("get-item failed"); + let item = { + #[cfg(not(Py_LIMITED_API))] + { + unsafe { list.get_item_unchecked(target_index) } + } + + #[cfg(Py_LIMITED_API)] + { + list.get_item(target_index).expect("get-item failed") + } + }; *length = Length(target_index); Some(item) } else { @@ -705,23 +682,11 @@ impl<'py> Iterator for BoundListIterator<'py> { } #[inline] - #[cfg(feature = "nightly")] + #[cfg(not(feature = "nightly"))] fn nth(&mut self, n: usize) -> Option { - #[cfg(not(Py_LIMITED_API))] - { - self.with_critical_section(|index, length, list| unsafe { - Self::nth_unchecked(index, length, list, n) - }) - } - #[cfg(Py_LIMITED_API)] - { - let Self { - index, - length, - list, - } = self; + self.with_critical_section(|index, length, list| { Self::nth(index, length, list, n) - } + }) } #[inline] @@ -851,7 +816,7 @@ impl<'py> Iterator for BoundListIterator<'py> { } #[inline] - #[cfg(all(not(Py_LIMITED_API), feature = "nightly"))] + #[cfg(feature = "nightly")] fn advance_by(&mut self, n: usize) -> Result<(), NonZero> { self.with_critical_section(|index, length, list| { let max_len = length.0.min(list.len()); @@ -898,23 +863,9 @@ impl DoubleEndedIterator for BoundListIterator<'_> { } #[inline] - #[cfg(feature = "nightly")] + #[cfg(not(feature = "nightly"))] fn nth_back(&mut self, n: usize) -> Option { - #[cfg(not(Py_LIMITED_API))] - { - self.with_critical_section(|index, length, list| unsafe { - Self::nth_back_unchecked(index, length, list, n) - }) - } - #[cfg(Py_LIMITED_API)] - { - let Self { - index, - length, - list, - } = self; - Self::nth_back(index, length, list, n) - } + self.with_critical_section(|index, length, list| Self::nth_back(index, length, list, n)) } #[inline] From 51104a1cd168ffc7593a22e8460961e963f3d2a9 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Tue, 14 Jan 2025 22:14:34 +0800 Subject: [PATCH 14/19] fix fmt --- src/types/list.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index 46ac31b5254..82b77ffc596 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -559,11 +559,13 @@ impl<'py> BoundListIterator<'py> { let target_index = index.0 + n; if target_index < length { let item = { - #[cfg(Py_LIMITED_API)] { + #[cfg(Py_LIMITED_API)] + { list.get_item(target_index).expect("get-item failed") } - #[cfg(not(Py_LIMITED_API))] { + #[cfg(not(Py_LIMITED_API))] + { unsafe { list.get_item_unchecked(target_index) } } }; @@ -684,9 +686,7 @@ impl<'py> Iterator for BoundListIterator<'py> { #[inline] #[cfg(not(feature = "nightly"))] fn nth(&mut self, n: usize) -> Option { - self.with_critical_section(|index, length, list| { - Self::nth(index, length, list, n) - }) + self.with_critical_section(|index, length, list| Self::nth(index, length, list, n)) } #[inline] From cae09816085c1ba405ecd1dc61cee252479031a5 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Tue, 14 Jan 2025 22:40:12 +0800 Subject: [PATCH 15/19] fix failing CI --- src/types/list.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index 82b77ffc596..5430731dbe5 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -548,7 +548,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(not(feature = "nightly"))] + #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] fn nth( index: &mut Index, length: &mut Length, @@ -619,7 +619,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(not(feature = "nightly"))] + #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] fn nth_back( index: &mut Index, length: &mut Length, @@ -684,7 +684,7 @@ impl<'py> Iterator for BoundListIterator<'py> { } #[inline] - #[cfg(not(feature = "nightly"))] + #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] fn nth(&mut self, n: usize) -> Option { self.with_critical_section(|index, length, list| Self::nth(index, length, list, n)) } @@ -863,7 +863,7 @@ impl DoubleEndedIterator for BoundListIterator<'_> { } #[inline] - #[cfg(not(feature = "nightly"))] + #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] fn nth_back(&mut self, n: usize) -> Option { self.with_critical_section(|index, length, list| Self::nth_back(index, length, list, n)) } From 00e48024d30fe2cb78aa7fcbc2ba08724b6c4d70 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Wed, 15 Jan 2025 22:33:02 +0800 Subject: [PATCH 16/19] Impl advance_back_by. Remove cfg flag for with_critical_section --- src/types/list.rs | 64 +++++++++++++++++++++++++++++++++-------------- 1 file changed, 45 insertions(+), 19 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index 5430731dbe5..a3071a22c17 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -647,7 +647,6 @@ impl<'py> BoundListIterator<'py> { } } - #[cfg(not(Py_LIMITED_API))] fn with_critical_section( &mut self, f: impl FnOnce(&mut Index, &mut Length, &Bound<'py, PyList>) -> R, @@ -819,25 +818,12 @@ impl<'py> Iterator for BoundListIterator<'py> { #[cfg(feature = "nightly")] fn advance_by(&mut self, n: usize) -> Result<(), NonZero> { self.with_critical_section(|index, length, list| { - let max_len = length.0.min(list.len()); - let currently_at = index.0; - if currently_at >= max_len { - if n == 0 { - return Ok(()); - } else { - return Err(unsafe { NonZero::new_unchecked(n) }); + for i in 0..n { + if unsafe { Self::next_unchecked(index, length, list).is_none() } { + return Err(unsafe { NonZero::new_unchecked(n - i) }); } } - - let items_left = max_len - currently_at; - if n <= items_left { - index.0 += n; - Ok(()) - } else { - index.0 = max_len; - let remainder = n - items_left; - Err(unsafe { NonZero::new_unchecked(remainder) }) - } + Ok(()) }) } } @@ -900,6 +886,19 @@ impl DoubleEndedIterator for BoundListIterator<'_> { R::from_output(accum) }) } + + #[inline] + #[cfg(feature = "nightly")] + fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero> { + self.with_critical_section(|index, length, list| { + for i in 0..n { + if unsafe { Self::next_back_unchecked(index, length, list).is_none() } { + return Err(unsafe { NonZero::new_unchecked(n - i) }); + } + } + Ok(()) + }) + } } impl ExactSizeIterator for BoundListIterator<'_> { @@ -1638,7 +1637,8 @@ mod tests { assert_eq!(iter.next().unwrap().extract::().unwrap(), 10); let mut iter = list.iter(); - iter.nth_back(1); + println!("iter.nth_back(1) = {:?}", iter.nth_back(1)); + // assert_eq!(iter.nth_back(1).unwrap().extract::().unwrap(), 9); assert_eq!(iter.nth(2).unwrap().extract::().unwrap(), 8); assert!(iter.next().is_none()); }); @@ -1726,4 +1726,30 @@ mod tests { assert_eq!(iter4.next().unwrap().extract::().unwrap(), 1); }) } + + #[cfg(feature = "nightly")] + #[test] + fn test_iter_advance_back_by() { + Python::with_gil(|py| { + let v = vec![1, 2, 3, 4, 5]; + let ob = (&v).into_pyobject(py).unwrap(); + let list = ob.downcast::().unwrap(); + + let mut iter = list.iter(); + assert_eq!(iter.advance_back_by(2), Ok(())); + assert_eq!(iter.next_back().unwrap().extract::().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 = list.iter(); + assert_eq!(iter2.advance_back_by(6), Err(NonZero::new(1).unwrap())); + + let mut iter3 = list.iter(); + assert_eq!(iter3.advance_back_by(5), Ok(())); + + let mut iter4 = list.iter(); + assert_eq!(iter4.advance_back_by(0), Ok(())); + assert_eq!(iter4.next_back().unwrap().extract::().unwrap(), 5); + }) + } } From b7373aaa145c1ee68e01c39273f658a4df9ff80d Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Thu, 16 Jan 2025 00:06:52 +0800 Subject: [PATCH 17/19] refactor advance_by and advance_back_by. Add back cfg for with_critical_section --- src/types/list.rs | 48 ++++++++++++++++++++++++++++++++++++----------- 1 file changed, 37 insertions(+), 11 deletions(-) diff --git a/src/types/list.rs b/src/types/list.rs index a3071a22c17..2d65dd4280f 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -640,13 +640,14 @@ impl<'py> BoundListIterator<'py> { list.get_item(target_index).expect("get-item failed") } }; - *length = Length(target_index); + length.0 = target_index; Some(item) } else { None } } + #[cfg(not(Py_LIMITED_API))] fn with_critical_section( &mut self, f: impl FnOnce(&mut Index, &mut Length, &Bound<'py, PyList>) -> R, @@ -818,12 +819,25 @@ impl<'py> Iterator for BoundListIterator<'py> { #[cfg(feature = "nightly")] fn advance_by(&mut self, n: usize) -> Result<(), NonZero> { self.with_critical_section(|index, length, list| { - for i in 0..n { - if unsafe { Self::next_unchecked(index, length, list).is_none() } { - return Err(unsafe { NonZero::new_unchecked(n - i) }); + let max_len = length.0.min(list.len()); + let currently_at = index.0; + if currently_at >= max_len { + if n == 0 { + return Ok(()); + } else { + return Err(unsafe { NonZero::new_unchecked(n) }); } } - Ok(()) + + let items_left = max_len - currently_at; + if n <= items_left { + index.0 += n; + Ok(()) + } else { + index.0 = max_len; + let remainder = n - items_left; + Err(unsafe { NonZero::new_unchecked(remainder) }) + } }) } } @@ -891,12 +905,25 @@ impl DoubleEndedIterator for BoundListIterator<'_> { #[cfg(feature = "nightly")] fn advance_back_by(&mut self, n: usize) -> Result<(), NonZero> { self.with_critical_section(|index, length, list| { - for i in 0..n { - if unsafe { Self::next_back_unchecked(index, length, list).is_none() } { - return Err(unsafe { NonZero::new_unchecked(n - i) }); + let max_len = length.0.min(list.len()); + let currently_at = index.0; + if currently_at >= max_len { + if n == 0 { + return Ok(()); + } else { + return Err(unsafe { NonZero::new_unchecked(n) }); } } - Ok(()) + + let items_left = max_len - currently_at; + if n <= items_left { + length.0 = max_len - n; + Ok(()) + } else { + length.0 = max_len; + let remainder = n - items_left; + Err(unsafe { NonZero::new_unchecked(remainder) }) + } }) } } @@ -1637,8 +1664,7 @@ mod tests { assert_eq!(iter.next().unwrap().extract::().unwrap(), 10); let mut iter = list.iter(); - println!("iter.nth_back(1) = {:?}", iter.nth_back(1)); - // assert_eq!(iter.nth_back(1).unwrap().extract::().unwrap(), 9); + assert_eq!(iter.nth_back(1).unwrap().extract::().unwrap(), 9); assert_eq!(iter.nth(2).unwrap().extract::().unwrap(), 8); assert!(iter.next().is_none()); }); From 7751a1c1844262a0be409ad1f692e9b1be2c4fc8 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Thu, 16 Jan 2025 00:30:04 +0800 Subject: [PATCH 18/19] Put allow deadcode for with_critical_section --- src/types/list.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/types/list.rs b/src/types/list.rs index 2d65dd4280f..155f5e57be8 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -647,7 +647,7 @@ impl<'py> BoundListIterator<'py> { } } - #[cfg(not(Py_LIMITED_API))] + #[allow(dead_code)] fn with_critical_section( &mut self, f: impl FnOnce(&mut Index, &mut Length, &Bound<'py, PyList>) -> R, From a735850cfedb52c02ac917f2b6433547155ac1f0 Mon Sep 17 00:00:00 2001 From: Owen Leung Date: Thu, 16 Jan 2025 19:28:05 +0800 Subject: [PATCH 19/19] Remove use of get_item. Revise changelog --- newsfragments/4810.added.md | 2 +- src/types/list.rs | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/newsfragments/4810.added.md b/newsfragments/4810.added.md index e89e22e544d..00c7c9e1127 100644 --- a/newsfragments/4810.added.md +++ b/newsfragments/4810.added.md @@ -1 +1 @@ -Optimizes `nth` and `nth_back` for `BoundListIterator` \ No newline at end of file +Optimizes `nth`, `nth_back`, `advance_by` and `advance_back_by` for `BoundListIterator` \ No newline at end of file diff --git a/src/types/list.rs b/src/types/list.rs index 155f5e57be8..51af830d6f5 100644 --- a/src/types/list.rs +++ b/src/types/list.rs @@ -548,7 +548,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] + #[cfg(not(feature = "nightly"))] fn nth( index: &mut Index, length: &mut Length, @@ -619,7 +619,7 @@ impl<'py> BoundListIterator<'py> { } #[inline] - #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] + #[cfg(not(feature = "nightly"))] fn nth_back( index: &mut Index, length: &mut Length, @@ -684,7 +684,7 @@ impl<'py> Iterator for BoundListIterator<'py> { } #[inline] - #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] + #[cfg(not(feature = "nightly"))] fn nth(&mut self, n: usize) -> Option { self.with_critical_section(|index, length, list| Self::nth(index, length, list, n)) } @@ -863,7 +863,7 @@ impl DoubleEndedIterator for BoundListIterator<'_> { } #[inline] - #[cfg(all(not(Py_LIMITED_API), not(feature = "nightly")))] + #[cfg(not(feature = "nightly"))] fn nth_back(&mut self, n: usize) -> Option { self.with_critical_section(|index, length, list| Self::nth_back(index, length, list, n)) }