Fix an edge case with StepBy::nth on non-fused iterators

[asder8215]

Fixes #159965.

From the nth documentation as pointed out by theemathas:

nth() will return None if n is greater than or equal to the length of the iterator.

Currently, there is an edge case for non-fused iterator where it's able to return Some through StepBy::nth iterator even though the first item from the non-fused iterator returns None (it is logically an empty iterator).

The issue came from how in the first take block it advances the underlying iterator forward using .next() and does not check if what it returns is a None value or not. This wouldn't be a problem for fused iterators because all the values that it would return after reaching the None point will also be None. However since non-fused iterators do not have to abide by continuously yielding None after reaching a None, it allows for a case, where after falling down from first_take block it can return Some(_) from a self.iter.nth() call in there.

In the first_take block we should definitely check if the first item we got from self.iter.next() is None item, and return None if it is so.

[rustbot]

r? @JohnTitor

rustbot has assigned @JohnTitor.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

Why was this reviewer chosen?

The reviewer was selected based on:

• Owners of files modified in this PR: libs
• libs expanded to 12 candidates
• Random selection from 7 candidates

[theemathas]

The later calls to self.iter.nth can also return None.

[asder8215]

The later calls to self.iter.nth can also return None.

Ah you're talking about when it overflows? You're right, I'll fix that when I'm back on my PC.

[asder8215]

@theemathas I guess the one thing that I have a question about with self.iter.nth() is: should we advance the iterator forward still after observing that self.iter.nth() returned None?

The reason why I ask that is because this would affect both non-fused/fused iterators at least from understanding this test case that I'm seeing here:

#[test]
#[allow(non_local_definitions)]
fn test_iterator_step_by_nth_overflow() {
    #[cfg(target_pointer_width = "16")]
    type Bigger = u32;
    #[cfg(target_pointer_width = "32")]
    type Bigger = u64;
    #[cfg(target_pointer_width = "64")]
    type Bigger = u128;

    #[derive(Clone)]
    struct Test(Bigger);
    impl Iterator for &mut Test {
        type Item = i32;
        fn next(&mut self) -> Option<Self::Item> {
            Some(21)
        }
        fn nth(&mut self, n: usize) -> Option<Self::Item> {
            self.0 += n as Bigger + 1;
            Some(42)
        }
    }

    let mut it = Test(0);
    let root = usize::MAX >> (usize::BITS / 2);
    let n = root + 20;
    (&mut it).step_by(n).nth(n);
    assert_eq!(it.0, n as Bigger * n as Bigger);

    // large step
    let mut it = Test(0);
    (&mut it).step_by(usize::MAX).nth(5);
    assert_eq!(it.0, (usize::MAX as Bigger) * 5);

    // n + 1 overflows
    let mut it = Test(0);
    (&mut it).step_by(2).nth(usize::MAX);
    assert_eq!(it.0, (usize::MAX as Bigger) * 2);

    // n + 1 overflows
    let mut it = Test(0);
    (&mut it).step_by(1).nth(usize::MAX);
    assert_eq!(it.0, (usize::MAX as Bigger) * 1);
}

The thing that I'm thinking about is if we have a custom nth function for our fused iterator that, let's say returns None always, and we create a StepBy iterator through a mutable borrow of the iterator, should our StepBy iterator stop calling self.iter.nth() in the first time it sees a None (in which case we'd only observe assert_eq!(it.0, n)) or should we still run through self.iter.nth() after finishing our overflow check (in which case we'd observe assert_eq!(it.0, n * n)).

My personal thought is the former because it wouldn't be advancing a non-fused iterator past the first None point it observed. However, this decision does affect behavior on both fused/non-fused.

[asder8215]

Out of caution, I'm going to mark this as libs-api nominated.

@rustbot label +I-libs-api-nominated

To clarify what the problem is, if you look into what StepBy::nth does, there's an overflow handling loop that occurs at the end:

default fn spec_nth(&mut self, mut n: usize) -> Option<I::Item> {
        ....
        // overflow handling
        loop {
            let mul = n.checked_mul(step);
            {
                if intrinsics::likely(mul.is_some()) {
                    return self.iter.nth(mul.unwrap() - 1);
                }
            }
            let div_n = usize::MAX / n;
            let div_step = usize::MAX / step;
            let nth_n = div_n * n;
            let nth_step = div_step * step;
            let nth = if nth_n > nth_step {
                step -= div_n;
                nth_n
            } else {
                n -= div_step;
                nth_step
            };
            self.iter.nth(nth - 1);
        }
}

This loop repeatedly calls on self.iter.nth(n - 1) until we reach a point where n * step doesn't overflow. The question here is how should we handle non-fused iterators in this scenario where they could return some Some(_) values, then return None, and then go back to returning Some(_) values (before we reach an overflow point/usize::MAX elements)? Should we stop the loop on the first sight of seeing a None value from self.iter.nth(nth - 1), or should we reserve returning None until after we fully resolve the overflow? If the former, this would mean that fused and non-fused iterators could have custom nth function be affected in the event that our StepBy iterator is constructed from a mutable borrow of the original iterator (see the test case I reference above). If the latter, we cause non-fused iterators to advance forward from the point that we receive a None from self.iter.nth(nth - 1) (which might not be the ideal behavior we intend for non-fused iterators).

[theemathas]

I don't understand what the problem here. We should return None as soon as self.iter.nth() returns None, right? I fail to see the issue with doing that.

I don't understand what your code snippet is supposed to demonstrate. Your description after the code talks about methods that always return None, but the iterator implementation in the code has methods that always returns Some.

[theemathas]

Here's a test case that hopefully demonstrates that we should stop iterating after any self.iter.nth() call returns None:

// An iterator that returns Some(0), then None,
// then returns Some(1) indefinitely
enum Wonky {
    First,
    Second,
    After,
}
impl Iterator for Wonky {
    type Item = i32;
    fn next(&mut self) -> Option<i32> {
        match self {
            Wonky::First => {
                *self = Wonky::Second;
                Some(0)
            },
            Wonky::Second => {
                *self = Wonky::After;
                None
            },
            Wonky::After => {
                Some(1)
            }
        }
    }
}

fn main() {
    let mut iter = Wonky::First.step_by(1);
    let _ = iter.next();
    // Currently returns Some(1) in release mode, which is wrong.
    // It should return None.
    // (It takes too long to run in debug mode.)
    println!("{:?}", iter.nth(usize::MAX));
}

[asder8215]

I don't understand what your #160025 (comment) is supposed to demonstrate. Your description after the code talks about methods that always return None, but the iterator implementation in the code has methods that always returns Some.

My bad, I referenced that test case because it made me realize that overflow could affect how the original iterator custom nth function is being called multiple times. To demonstrate the issue:

#[test]
#[allow(non_local_definitions)]
fn test_iterator_step_by_nth_overflow() {
    type Bigger = u128;

    #[derive(Clone)]
    struct Test(Bigger);
    impl Iterator for &mut Test {
        type Item = i32;
        fn next(&mut self) -> Option<Self::Item> {
            None
        }
        fn nth(&mut self, n: usize) -> Option<Self::Item> {
            self.0 += n as Bigger + 1;
            None
        }
    }

    let mut it = Test(0);
    let root = usize::MAX >> (usize::BITS / 2);
    let n = root + 20;
    (&mut it).step_by(n).nth(n);
    assert_eq!(it.0, n as Bigger * n as Bigger); // our iterator u128 value would hold n * n
}

If we go by what the current code does it doesn't handle returning None early on fused iterator on overflow (only until after we resolve the overflow n * step). That also means that the underlying iterator custom nth function is called n times, which from the test case above would have the Test struct wrap around the value n * n.

However, if we decide to return None early on the first instance of seeing it from self.iter.nth(), we'd observe that that this Test struct would wrap around the value n.

I get that for non-fused iterator that we should return None on the first sight that we see None from self.iter.nth(). That makes sense to me so that we don't advance iterator forward to a section where we observe Some() values. But does that mean we consider the behavior with fused iterators calling on the original iterator custom nth function n times a bug on overflow or should that be intended behavior?

[theemathas]

But does that mean we consider the behavior with fused iterators calling on the original iterator custom nth function n times a bug on overflow or should that be intended behavior?

I personally consider this previous behavior to either be a bug, or be an internal implementation detail that's not stably guaranteed.

[nia-e]

Per today's meeting, we think this is not necessarily a libs-api thing and behaviour is up to reviewer discretion; we make no stable guarantees as to what StepBy means for non-fused iterators. Either way, thanks for checking ^^

[asder8215]

That sounds good to me, I'll push the change to early return on None from self.iter.nth(). Will leave it to JohnTitor on if the behavior is okay.

[JohnTitor]

@rustbot reroll

[clarfonthey]

Perhaps using the ? try syntax would be cleaner? This would be return Some(first) then, but in the other places this kind of branch is used it would be cleaner.

View changes since the review

[clarfonthey]

@rustbot author

[rustbot]

Reminder, once the PR becomes ready for a review, use @rustbot ready.

[asder8215]

@rustbot ready

[clarfonthey]

Latest version looks okay to me. Thank you!

@bors r+ rollup

[rust-bors]

📌 Commit 366ebd6 has been approved by clarfonthey

It is now in the queue for this repository.