How to implement IntoIterator using FromFn?

I think this MRE gives an idea of what I am trying to achieve. I am a bit stuck on (a) what the type IntoIter in the trait should be and (b) whether it is necessary to have a double closure inside from_fn...


#[derive(Default)]
struct Counter {
    count: u32
}

impl IntoIterator for Counter {
    type Item = u32;

    type IntoIter = std::iter::FromFn<FnMut() -> Option<Self::Item>>;

    fn into_iter(self) -> Self::IntoIter {
        std::iter::from_fn(move || {
            // move self into the outer closure
            let mut counter = self;
            // this is the closure that will be continously produce Self::Item?
            || {
                counter += 1;
                Some(counter)
            }
        })
    }
}

fn main() {
    eprintln!("{:?}", Counter::default().into_iter().take(5).collect::<Vec<_>>())
}

There’s multiple options with trade-off between convenience/performance here. The optimal solution would probably be to use type_alias_impl_traits but that’s not a stable feature of Rust yet. Btw. the double closure thing doesn’t make sense and isn’t necessary. With that unstable feature it would look like

#![feature(type_alias_impl_trait)]

#[derive(Default)]
struct Counter {
    count: u32,
}

impl IntoIterator for Counter {
    type Item = u32;

    type IntoIter = std::iter::FromFn<impl FnMut() -> Option<Self::Item>>;

    fn into_iter(mut self) -> Self::IntoIter {
        std::iter::from_fn(move || {
            self.count += 1;
            Some(self.count)
        })
    }
}

fn main() {
    eprintln!(
        "{:?}",
        Counter::default().into_iter().take(5).collect::<Vec<_>>()
    )
}

Rust Playground

or

#![feature(type_alias_impl_trait)]

#[derive(Default)]
struct Counter {
    count: u32,
}

impl IntoIterator for Counter {
    type Item = u32;

    type IntoIter = impl Iterator<Item = Self::Item>;

    fn into_iter(mut self) -> Self::IntoIter {
        std::iter::from_fn(move || {
            self.count += 1;
            Some(self.count)
        })
    }
}

fn main() {
    eprintln!(
        "{:?}",
        Counter::default().into_iter().take(5).collect::<Vec<_>>()
    )
}

Rust Playground


On stable Rust, you can do a similar approach (but using dynamic function calls) by using trait objects instead, and you’ll have to box them:

#[derive(Default)]
struct Counter {
    count: u32,
}

impl IntoIterator for Counter {
    type Item = u32;

    type IntoIter = std::iter::FromFn<Box<dyn FnMut() -> Option<Self::Item>>>;

    fn into_iter(mut self) -> Self::IntoIter {
        std::iter::from_fn(Box::new(move || {
            self.count += 1;
            Some(self.count)
        }))
    }
}

fn main() {
    eprintln!(
        "{:?}",
        Counter::default().into_iter().take(5).collect::<Vec<_>>()
    )
}

Rust Playground

or

#[derive(Default)]
struct Counter {
    count: u32,
}

impl IntoIterator for Counter {
    type Item = u32;

    type IntoIter = Box<dyn Iterator<Item = Self::Item>>;

    fn into_iter(mut self) -> Self::IntoIter {
        Box::new(std::iter::from_fn(move || {
            self.count += 1;
            Some(self.count)
        }))
    }
}

fn main() {
    eprintln!(
        "{:?}",
        Counter::default().into_iter().take(5).collect::<Vec<_>>()
    )
}

Rust Playground


Regarding the question whether the Iterator or FnMut is the right abstraction level, I suppose I’d prefer the Iterator (for either the unstable impl Iterator<…> or the stable Box<dyn Iterator<…> approach. Depending on the use case, for the trait object it might also make sense (or even be important/necesary) to include some auto-trait bounds, so something like Box<dyn Iterator<Item = Self::Item> + Send + Sync>.


Now, the alternative (and slightly more verbose) approach is to work with a manual Iterator implementation instead. You’d define a struct that captures the necessary state, instead of the closure, and you’d implement Iterator for that. This avoids any dynamic function calls, so it’s more performant; the main downside is that it might be somewhat more verbose to type, but decide for yourself.

#[derive(Default)]
struct Counter {
    count: u32,
}

impl IntoIterator for Counter {
    type Item = u32;

    type IntoIter = IntoIter;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter(self)
    }
}

// Naming-wise, if this whole thing lives in a module called `counter`, a generic
// "IntoIter" name makes sense, similar to how many standard library collections do it.
struct IntoIter(Counter);
impl Iterator for IntoIter {
    type Item = u32;
    fn next(&mut self) -> Option<Self::Item> {
        self.0.count += 1;
        Some(self.0.count)
    }
}

fn main() {
    eprintln!(
        "{:?}",
        Counter::default().into_iter().take(5).collect::<Vec<_>>()
    )
}

Rust Playground


For this particular use-case of an infinite iterator note that you could use repeat_with in place of from_fn, too, which does e.g. give a more precise size_hint implementation. For the manual Iterator you could, similarly, implement size_hint accordingly, in this example returning (usize::MAX, None).

2 Likes

Thank you for such a comprehensive answer with so many alternatives! I had considered repeat_with but the problem there would have been that there is no way to terminate it, which is why I ended up working with from_fn.

Of course, that’s kind of the point of it :slight_smile: . I assume that your non-MRE use-case involves an iterator that will eventually terminate, and in this case repeat_with no longer works; it was just this concrete code example of yours that unconditionally returned Some(…), and I wanted to point it out in case you weren’t aware of repeat_with :wink:

1 Like

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.