How to avoid unwrap on endless iterators?

help
1

Some iterators are known to never be exhausted, such as std::iter::Cycle as returned by Iterator::cycle. Yet, we need to call Option::unwrap on each value returned by Iterator::next. We know unwrapping will always succeed, so we might as well use Option::unwrap_unchecked, but that would require using unsafe code.

It would be nicer if the unsafe code was hidden in some generic code instead, so I came up with this:

unsafe trait EndlessIterator: Iterator {
    fn next_unwrapped(&mut self) -> Self::Item {
        let item = Iterator::next(self);
        unsafe { item.unwrap_unchecked() }
    }
}

unsafe impl<I> EndlessIterator for std::iter::Cycle<I>
where
    I: Clone + Iterator,
{
}

fn main() {
    let vec = vec![1, 2, 3];
    let mut iter = vec.iter().cycle();
    for _ in 0..10 {
        println!("{}", iter.next_unwrapped());
    }
}

(Playground)

My questions:

  1. Is this sound?
  2. Does such a trait or method exist already?
2

It's pretty rare to need to call next() manually. What's your use case? Can't you just replace calls to next() with something that already encapsulates (or doesn't care about) the unwrapping, e.g. a for loop, iterator adaptors, or something like that?

3

Multiplying complex numbers (that I receive in slices) with pre-calculated numbers.

pub async fn freq_shift(
    mut input: Receiver<Complex<f32>>,
    output: Sender<Complex<f32>>,
    numerator: isize,
    denominator: usize,
) {
    const TAU: f32 = std::f32::consts::TAU;
    let mut phase_vec: Vec<Complex<f32>> = Vec::new();
    for i in 0..denominator {
        let (im, re) = <f32>::sin_cos(i as f32 * numerator as f32 / denominator as f32 * TAU);
        let factor = Complex::new(re, im);
        phase_vec.push(factor);
    }
    let mut phase_iter = phase_vec.iter().cycle();
    let mut next_phase = || phase_iter.next().unwrap(); // here is the .unwrap() that could be .unwrap_unchecked()
    let mut buf_pool = ChunkBufPool::<Complex<f32>>::new();
    loop {
        match input.recv().await {
            Ok(input_chunk) => {
                let mut output_chunk = buf_pool.get();
                for sample in input_chunk.iter() {
                    output_chunk.push(sample * next_phase());
                }
                output.send(output_chunk.finalize());
            }
            Err(RecvError::Closed) => return,
            Err(RecvError::Lagged) => output.lag(),
        }
    }
}

Likely it could be optimized somehow. I guess this approach isn't as fast as it could be.

Note that input_chunk is a smart-pointer to an immutable [Complex<f32>] (contains an Arc<Complex<f32>>), so it's no ordinary slice.

4

Definitely not; for one thing, std::iter::Cycle is not (always) endless.

To guarantee that Cycle<I> is endless, I has to be nonempty and have no interior mutability. It would be doable in today's Rust for some iterators but you'd need dependent types in your case (to prove that the std::slice::Iter is always nonempty).

4 Likes
5

Another simple example is just std::iter::empty().cycle()

6

I assume that the latter could be fixed with this?

-unsafe impl<I> EndlessIterator for std::iter::Cycle<I>
-where
-    I: Clone + Iterator,
+unsafe impl<'a, T> EndlessIterator for std::iter::Cycle<std::slice::Iter<'a, T>>

But this wouldn't fix the empty slice case: Playground with UB.

This could work:

#[derive(Clone)]
pub struct NonEmptySliceIter<'a, T>(std::slice::Iter<'a, T>);

impl<'a, T> Iterator for NonEmptySliceIter<'a, T> {
    type Item = &'a T;
    fn next(&mut self) -> Option<&'a T> {
        self.0.next()
    }
}

impl<'a, T> NonEmptySliceIter<'a, T> {
    fn new(slice: &'a [T]) -> Self {
        assert!(slice.len() > 0);
        Self(slice.iter())
    }
}

pub unsafe trait EndlessIterator: Iterator {
    fn next_unwrapped(&mut self) -> Self::Item {
        let item = Iterator::next(self);
        unsafe { item.unwrap_unchecked() }
    }
}

unsafe impl<'a, T> EndlessIterator for std::iter::Cycle<NonEmptySliceIter<'a, T>>
where
    T: Clone,
{
}

fn main() {
    let vec = vec![1, 2, 3]; // try to make this empty
    let mut iter = NonEmptySliceIter::new(&vec).cycle();
    for _ in 0..10 {
        println!("{}", iter.next_unwrapped());
    }
}

(Playground)

But feels pretty bloated.

7

Is the following sound?

pub fn cycle<'a, T>(slice: &'a [T]) -> impl FnMut() -> &'a T {
    assert!(slice.len() > 0);
    let mut iter = slice.iter().cycle();
    move || {
        let option = iter.next();
        unsafe { option.unwrap_unchecked() }
    }
}

fn main() {
    let vec = vec![1, 2, 3];
    let mut next = cycle(&vec);
    for _ in 0..10 {
        println!("{}", next());
    }
}

(Playground)

8

Your function looks sound. The assertion is sufficient to prevent the initial slice::Iter from returning None immediately.

I doesn't need interior mutability, if it has an evil Clone implementation (Rust Playground):

struct EvilIter(bool);

impl Clone for EvilIter {
    fn clone(&self) -> Self {
        Self(false)
    }
}

impl Iterator for EvilIter {
    type Item = &'static str;

    fn next(&mut self) -> Option<Self::Item> {
        std::mem::take(&mut self.0).then_some("Hello, world!")
    }
}

fn main() {
    let mut iter = EvilIter(true);
    assert_eq!(iter.next(), Some("Hello, world!")); // non-empty
    assert_eq!(iter.next(), None);
    let mut iter = EvilIter(true).cycle();
    assert_eq!(iter.next(), Some("Hello, world!"));
    assert_eq!(iter.next(), None);
}
5 Likes
9

This is a good example of a place where the compiler will almost certainly optimize away the check in the safe code, because if the iterator is really endless, then the code "obviously" only ever returns Some(…).

Here's a simple demonstration: https://rust.godbolt.org/z/4E9xMsoab

pub fn demo(it: &mut std::iter::Repeat<i32>) -> i32 {
    it.next().unwrap()
}

That compiles down to just

example::demo:
        mov     eax, dword ptr [rdi]
        ret

That's no checks, no panicking paths. Just a copy of the integer because it's obvious to the optimizer that that's what it always does.

So trying to prove that the unsafe is actually correct here is error-prone -- as the Cycle example demonstrates -- and usually not actually a speed win. Thus it's probably not worth bothering.

Also, if you want a trait for this, it might be worth getting it added to core, as was recently discussed on IRLO: Infinite iterators in for - #2 by scottmcm - language design - Rust Internals, since then it would be up to core to forward it correctly through adapters and such.

9 Likes
10

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.