How to implement an infinite integer iterator?

Question

I have created an iterator to generate an infinite list of integers from 0 onwards as shown below.

struct Infinite {
    current: i64,
}

impl Infinite {
    fn new() -> Infinite {
        Infinite { current: -1 }
    }
}

impl Iterator for Infinite {
    type Item = i64;

    fn next(&mut self) -> Option<Self::Item> {
        self.current += 1;
        Some(self.current)
    }
}

I hope to manipulate this list of integers with different commands such as Multiply (to multiple each integer by a number n) or subtract (to subtract n from each integer). Thus, I created a Multiply struct and implemented an iterator for it as shown below.

struct Multiply {
    n: i64,
    inner: Box<dyn Iterator<Item = i64>>,
}

impl Multiply {
    fn new(inner: Box<dyn Iterator<Item = i64>>, n: i64) -> Multiply {
        Multiply { n, inner }
    }
}

impl Iterator for Multiply {
    type Item = i64;

    fn next(&mut self) -> Option<Self::Item> {
        Some(self.inner.next().unwrap() * self.n)
    }
}

However, I am not sure how to implement this in the main function. Should I begin by defining a new iter object this way?

let mut iter = Infinite::new();
    iter.next();

Then, how should I implement the Multiply iterator to iter ?

Answer 1

fn main() {
    let mut iter = Multiply::new(Box::new(Infinite::new()), 2);

    println!("{}", iter.next().unwrap());
    println!("{}", iter.next().unwrap());
    println!("{}", iter.next().unwrap());
}

0
2
4

---

That said, it's pretty unusual for an iterator to carry a Box object. It's usually done with a generic. (Also, don't use .unwrap() when mapping from one iterator to another, use Option::map() instead to preserve the end of the iterator)

struct Infinite {
    current: i64,
}

impl Infinite {
    fn new() -> Infinite {
        Infinite { current: -1 }
    }
}

impl Iterator for Infinite {
    type Item = i64;

    fn next(&mut self) -> Option<Self::Item> {
        self.current += 1;
        Some(self.current)
    }
}

struct Multiply<I> {
    n: i64,
    inner: I,
}

impl<I> Multiply<I> {
    fn new(inner: I, n: i64) -> Multiply<I> {
        Multiply { n, inner }
    }
}

impl<I> Iterator for Multiply<I>
where
    I: Iterator<Item = i64>,
{
    type Item = i64;

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next().map(|n| n * self.n)
    }
}

fn main() {
    let mut iter = Multiply::new(Infinite::new(), 2);

    println!("{}", iter.next().unwrap());
    println!("{}", iter.next().unwrap());
    println!("{}", iter.next().unwrap());
}

0
2
4

---

Further, this is usually done with an extension trait for Iterator so you are able to chain it. This is also how the entire itertools crate works.

struct Infinite {
    current: i64,
}

impl Infinite {
    fn new() -> Infinite {
        Infinite { current: -1 }
    }
}

impl Iterator for Infinite {
    type Item = i64;

    fn next(&mut self) -> Option<Self::Item> {
        self.current += 1;
        Some(self.current)
    }
}

struct Multiply<I>
where
    I: Iterator<Item = i64>,
{
    n: i64,
    inner: I,
}

impl<I> Iterator for Multiply<I>
where
    I: Iterator<Item = i64>,
{
    type Item = i64;

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next().map(|n| n * self.n)
    }
}

trait MultiplyExt
where
    Self: Iterator + Sized,
{
    fn multiply_with(self, n: i64) -> Multiply<Self>
    where
        Self: Iterator<Item = i64>,
    {
        Multiply { n, inner: self }
    }
}

impl<T: Iterator> MultiplyExt for T {}

fn main() {
    let mut iter = Infinite::new().multiply_with(2);

    println!("{}", iter.next().unwrap());
    println!("{}", iter.next().unwrap());
    println!("{}", iter.next().unwrap());
}

0
2
4