Optimize nth and nth_back for BoundListIterator #4810

Owen-CH-Leung · 2024-12-21T06:48:36Z

This PR optimizes nth and nth_back for BoundListIterator, and added unittest & benchmarks for these 2 APIs. Here are the benchmark of the optimized nth and nth_back:

On Stable toolchain

With Optimization

list_nth                time:   [351.21 ns 351.71 ns 352.28 ns]
list_nth_back           time:   [355.34 ns 356.09 ns 356.95 ns]

The default nth and nth_back implementation

list_nth                time:   [2.8422 µs 2.8502 µs 2.8592 µs]
list_nth_back           time:   [3.6509 µs 3.6557 µs 3.6611 µs]

On nightly toolchain

With Optimization

list_nth                time:   [293.11 ns 293.59 ns 294.15 ns]
list_nth_back           time:   [301.09 ns 301.51 ns 302.00 ns]

The default nth and nth_back implementation

list_nth                time:   [2.8488 µs 2.8525 µs 2.8568 µs]
list_nth_back           time:   [3.6524 µs 3.6600 µs 3.6683 µs]

…nchmarks

davidhewitt · 2024-12-21T09:10:52Z

Awesome, thanks! Will seek to review soon. I guess we can do the same thing for tuple iteration?

Owen-CH-Leung · 2024-12-21T15:48:58Z

Awesome, thanks! Will seek to review soon. I guess we can do the same thing for tuple iteration?

Yup - I'll file a PR to optimise tuple iterator after the holiday =)

davidhewitt · 2024-12-24T11:04:55Z

src/types/list.rs

+        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) };


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...

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.

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

ngoldbaum · 2025-01-06T19:00:30Z

Thanks for working on this! Hopefully my PR should be merged in the next few days and then you can rebase this.

ngoldbaum · 2025-01-06T19:02:40Z

It might also make sense to implement advance_by on nightly. nth is implemented in terms of advance_by, so this improvement would fall out "for free" on nightly if you implement that function.

ngoldbaum

See inline comments, I think there are some issues in how this is set up for the limited API. I also think we should define nth and nth_back for all non-nightly builds. I also don't think it's necessary to call get_item in advance_by, per my reading of the advance_by docs.

src/types/list.rs

pyo3-benches/benches/bench_list.rs

noxfile.py

Owen-CH-Leung · 2025-01-10T13:36:53Z

@ngoldbaum Can I seek your review again ? I've adopted your suggested changes above

the CI clippy/x86_64-unknown-linux-gnu/beta is failing as main btw

ngoldbaum · 2025-01-10T17:02:31Z

the CI clippy/x86_64-unknown-linux-gnu/beta is failing as main btw

GitHub's interface doesn't help by putting a big red x next to it, but this is fine and we can merge PRs as long as the "conclusion" job finishes. Right now it's broken because of an upstream bug in clippy, so it will probably stay that way until the beta channel gets the fix.

I'll try to give this PR another once-over today but failing that will get to it next week.

ngoldbaum · 2025-01-13T18:13:34Z

src/types/list.rs

@@ -750,6 +849,32 @@ impl<'py> Iterator for BoundListIterator<'py> {
            None
        })
    }
+
+    #[inline]
+    #[cfg(all(not(Py_LIMITED_API), feature = "nightly"))]


Suggested change

#[cfg(all(not(Py_LIMITED_API), feature = "nightly"))]

#[cfg(feature = "nightly")]

Unless I'm missing something...

I feel we need the flag not(Py_LIMITED_API. Otherwise we will compile the code without access to with_critical_section

https://github.com/PyO3/pyo3/actions/runs/12769998628/job/35593851381?pr=4810

I think you can delete the #[cfg(not(Py_LIMITED_API))] on BoundListIterator::with_critical_section, it was only there to avoid a clippy lint about unused code. pyo3::sync::with_critical_section is unconditionally exposed in the PyO3 API, it's just a no-op on GIL-enabled builds. We wanted to allow people to write code that avoids conditional compilation, where possible.

Thanks. I deleted #[cfg(not(Py_LIMITED_API))] but was hitting the unused code lint warning, so I put in a allow lint dead_code. Lemme know your thoughts.

I also pushed a few commits and implemented advance_back_by.

ngoldbaum · 2025-01-13T18:19:39Z

src/types/list.rs

@@ -547,6 +547,46 @@ impl<'py> BoundListIterator<'py> {
        }
    }

+    #[inline]
+    #[cfg(all(not(Py_LIMITED_API), feature = "nightly"))]


I'm confused - shouldn't this only be defined on not(feature = "nightly")? Because on nightly this is defined in terms of advance_by in the standard library implementation, so we don't need to override nth. It'll just do the right thing on nightly by calling advance_by and then next(). On not nightly you need this override because there's no way to override advance_by.

Ditto for all the other BoundListIterator methods.

Indeed - I spent hours reading the stdlib again to figure out that on stable toolchain, we should override nth, and override advance_by on nightly. Thanks for spotting this. I've made the changes to compile it on stable

ngoldbaum · 2025-01-13T18:21:50Z

src/types/list.rs

+
+    #[inline]
+    #[cfg(all(Py_LIMITED_API, feature = "nightly"))]
+    #[deny(unsafe_op_in_unsafe_fn)]


Suggested change

#[deny(unsafe_op_in_unsafe_fn)]

there's no unsafe code in this function so this is unnecessary

ngoldbaum · 2025-01-13T18:23:48Z

src/types/list.rs

+    ) -> Option<Bound<'py, PyAny>> {
+        let length = length.0.min(list.len());
+        let target_index = index.0 + n;
+        if index.0 + n < length {


Suggested change

if index.0 + n < length {

if target_index < length {

Thanks & revised

ngoldbaum · 2025-01-13T18:25:27Z

src/types/list.rs

+    ) -> Option<Bound<'py, PyAny>> {
+        let length = length.0.min(list.len());
+        let target_index = index.0 + n;
+        if index.0 + n < length {


Suggested change

if index.0 + n < length {

if target_index < length {

Thanks & revised

src/types/list.rs

ngoldbaum · 2025-01-13T18:37:06Z

src/types/list.rs

@@ -625,6 +704,26 @@ impl<'py> Iterator for BoundListIterator<'py> {
        }
    }

+    #[inline]
+    #[cfg(feature = "nightly")]


Suggested change

#[cfg(feature = "nightly")]

#[cfg(not(feature = "nightly"))]

Thanks & revised

ngoldbaum · 2025-01-13T18:40:35Z

src/types/list.rs

+        #[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)
+        }


Suggested change

#[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)

}

self.with_critical_section(|index, length, list| unsafe {

Self::nth(index, length, list, n)

})

If you implement my suggestion to only have BoundListIterator::nth then you can simplify this a lot.

and you can do a similar refactor for the nth_back implementation for the impl DoubleEndedIterator for BoundListIterator block below

Thanks. I've adopted this idea for nth and nth_back

ngoldbaum · 2025-01-13T18:43:59Z

src/types/list.rs

+    ) -> Option<Bound<'py, PyAny>> {
+        let length_size = length.0.min(list.len());
+        if index.0 + n < length_size {
+            let target_index = length_size - n - 1;


Personally this logic feels a little backwards to me, and I would probably prefer to do this using signed integers and testing if the index is less than zero. That said, this is totally equivalent and if it makes sense to you as-is then no need to change it.

… for proper compilation

Owen-CH-Leung · 2025-01-14T15:00:04Z

Just realize that I forgot to implement advance_back_by. I'll implement it in coming days

…al_section

ngoldbaum

Nice, much simpler now!

I'm not going to merge immediately because I'd like someone with a little bit more experience maintaining PyO3 to give this a once-over.

ngoldbaum · 2025-01-15T20:22:44Z

newsfragments/4810.added.md

@@ -0,0 +1 @@
+Optimizes `nth` and `nth_back` for `BoundListIterator`


maybe also mention that advance_by and advance_by_back are implemented on nightly

src/types/list.rs

davidhewitt

Thanks, this looks great to me! Sorry for the slow review cycles 😬

* Optimize nth and nth_back for BoundListIterator. Add unit test and benchmarks * Fix fmt and newsfragment CI * Fix clippy and changelog CI * Revise Impl of nth and nth_back. Impl advance_by * Fix failing fmt * Fix failing ruff test * branch out nth, nth_unchecked, nth_back, nth_back_unchecked. * Fix fmt * Revise advance_by impl. add advance_by unittest. * Fix fmt * Fix clippy unused function warning * Set appropriate Py_LIMITED_API flag * Rewrite nth & nth_back using conditional compilation. Rearrange flags for proper compilation * fix fmt * fix failing CI * Impl advance_back_by. Remove cfg flag for with_critical_section * refactor advance_by and advance_back_by. Add back cfg for with_critical_section * Put allow deadcode for with_critical_section * Remove use of get_item. Revise changelog

Owen-CH-Leung added 3 commits December 21, 2024 14:43

Optimize nth and nth_back for BoundListIterator. Add unit test and be…

cc9cabd

…nchmarks

Fix fmt and newsfragment CI

3a0c196

Fix clippy and changelog CI

40d38f3

Owen-CH-Leung marked this pull request as ready for review December 21, 2024 15:47

davidhewitt reviewed Dec 24, 2024

View reviewed changes

davidhewitt mentioned this pull request Jan 3, 2025

Implement locked iteration for PyList #4789

Merged

Owen-CH-Leung added 4 commits January 9, 2025 14:50

Merge branch 'main' into optimize_nth_and_nthback_for_boundlistiter

1b19616

Revise Impl of nth and nth_back. Impl advance_by

f6e95a8

Fix failing fmt

b0c749b

Fix failing ruff test

b2bf973

ngoldbaum reviewed Jan 9, 2025

View reviewed changes

Owen-CH-Leung mentioned this pull request Jan 9, 2025

ci: updates for Rust 1.84 #4846

Merged

Owen-CH-Leung added 3 commits January 10, 2025 08:52

Merge branch 'main' into optimize_nth_and_nthback_for_boundlistiter

4e86709

branch out nth, nth_unchecked, nth_back, nth_back_unchecked.

e4269c2

Fix fmt

6e18229

Owen-CH-Leung marked this pull request as draft January 10, 2025 08:22

Owen-CH-Leung added 4 commits January 10, 2025 18:53

Revise advance_by impl. add advance_by unittest.

5bab05b

Fix fmt

0b23173

Fix clippy unused function warning

e88f8be

Set appropriate Py_LIMITED_API flag

3a7a171

Owen-CH-Leung marked this pull request as ready for review January 10, 2025 13:35

Owen-CH-Leung requested a review from ngoldbaum January 10, 2025 13:37

ngoldbaum reviewed Jan 13, 2025

View reviewed changes

Rewrite nth & nth_back using conditional compilation. Rearrange flags…

ed8dba6

… for proper compilation

Owen-CH-Leung added 2 commits January 14, 2025 22:14

fix fmt

51104a1

fix failing CI

cae0981

Owen-CH-Leung added 3 commits January 15, 2025 22:33

Impl advance_back_by. Remove cfg flag for with_critical_section

00e4802

refactor advance_by and advance_back_by. Add back cfg for with_critic…

b7373aa

…al_section

Put allow deadcode for with_critical_section

7751a1c

ngoldbaum approved these changes Jan 15, 2025

View reviewed changes

Icxolu reviewed Jan 15, 2025

View reviewed changes

src/types/list.rs Show resolved Hide resolved

Remove use of get_item. Revise changelog

a735850

bschoenmaeckers mentioned this pull request Jan 29, 2025

Optimize Iterator::last() for PyList & PyTuple and Iterator::count() for PyDict, PyList, PyTuple & PySet #4878

Merged

Owen-CH-Leung requested review from Icxolu and davidhewitt February 2, 2025 07:05

davidhewitt approved these changes Feb 7, 2025

View reviewed changes

davidhewitt added this pull request to the merge queue Feb 7, 2025

Merged via the queue into PyO3:main with commit 903afcd Feb 7, 2025
45 of 46 checks passed

	#[cfg(all(not(Py_LIMITED_API), feature = "nightly"))]
	#[cfg(feature = "nightly")]

		@@ -0,0 +1 @@
		Optimizes `nth` and `nth_back` for `BoundListIterator`

Optimize nth and nth_back for BoundListIterator #4810

Optimize nth and nth_back for BoundListIterator #4810

Uh oh!

Conversation

Owen-CH-Leung commented Dec 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

On Stable toolchain

On nightly toolchain

Uh oh!

davidhewitt commented Dec 21, 2024

Uh oh!

Owen-CH-Leung commented Dec 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Owen-CH-Leung Dec 29, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Owen-CH-Leung Dec 29, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ngoldbaum commented Jan 6, 2025

Uh oh!

ngoldbaum commented Jan 6, 2025

Uh oh!

ngoldbaum left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Owen-CH-Leung commented Jan 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ngoldbaum commented Jan 10, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Owen-CH-Leung commented Jan 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Owen-CH-Leung commented Dec 21, 2024 •

edited

Loading

Owen-CH-Leung commented Dec 21, 2024 •

edited

Loading

Owen-CH-Leung Dec 29, 2024 •

edited

Loading

Owen-CH-Leung Dec 29, 2024 •

edited

Loading

Owen-CH-Leung commented Jan 10, 2025 •

edited

Loading

Owen-CH-Leung commented Jan 14, 2025 •

edited

Loading