How do you implement an Iterator whose successor depends on the index of the term?

Question

I'm trying to implement an Iterator that yields the sequence x^2, x^2+x, x^2+2x, x^2+3x... for a constant (at call time) argument x, in Rust.

As I understand it, all I have to work with at any point in the implementation is self.curr and self.next. It seems to me that the sequence depends on the index of the item in the sequence.

struct SquareMultiple {
    // generates the sequence j = i^2, i^2+i, i^2+2i, i^2+3i, ...,
    curr: i64,
    next: i64,
}

// Implement `Iterator` for `SquareMultiple`.
// The `Iterator` trait only requires a method to be defined for the `next`
// element.
impl Iterator for SquareMultiple {
    type Item = i64;

    // Here, we define the sequence using `.curr` and `.next`.
    // The return type is `Option<T>`:
    //     * When the `Iterator` is finished, `None` is returned.
    //     * Otherwise, the next value is wrapped in `Some` and returned.
    fn next(&mut self) -> Option<I64> {

        // FIXME: What to do here?
        let new_next = self.curr + self.next;

        self.curr = self.next;
        self.next = new_next;

        // Since there's no endpoint to a SquareMultiple sequence, the
        // `Iterator` will never return `None`, and `Some` is always returned.
        Some(self.curr)
    }
}

// Returns a SquareMultiple sequence generator
fn squareseq() -> SquareMultiple {
    SquareMultiple { curr: 1, next: 2 }
}

I also considered overloading the struct with a index attribute, but that seems somehow like an abuse of this pattern.

What's the Rustic way to go about this?

Answer 1

Another relatively simple variant using from_fn is missing here:

fn square_multiple(x: i64) -> impl Iterator<Item = i64> {
    let x2 = x * x;
    let mut v = -x;
    std::iter::from_fn(move || {
        v += x;
        Some(x2 + v)
    })
}

Playground

Answer 2

Rust 1.34

With iter::successors and impl trait, this can be greatly simplified:

fn square_multiple(x: i64) -> impl Iterator<Item = i64> {
    std::iter::successors(Some(x * x), move |v| Some(v + x))
}

Previously

You can store whatever you want in a struct that implements Iterator. It seems the simplest to me to store the current value and the value to increment by.

struct SquareMultiple {
    curr: i64,
    inc: i64,
}

impl Iterator for SquareMultiple {
    type Item = i64;

    fn next(&mut self) -> Option<i64> {
        let val = self.curr;
        self.curr += self.inc;
        Some(val)
    }
}

impl SquareMultiple {
    fn new(x: i64) -> Self {
        SquareMultiple { curr: x * x, inc: x }
    }
}

fn main() {
    for i in SquareMultiple::new(5).take(10) {
        println!("{}", i);
    }
}

It's probably worth documenting that the iterator goes forever and thus will either panic or wraparound when it exceeds 2^63.

I like this solution because it doesn't multiply at all. For some reason, my brain thinks that adding is "easier" than multiplying.

---

If you really need to use the index, use a RangeFrom and map it:

fn main() {
    let x = 5;
    let clever = (0..).map(|i| x * (x + i));
    for i in clever.take(10) {
        println!("{}", i);
    }
}

If you need a separate function and maximum performance (add appropriate sound effect), you can create a new type:

use std::ops::RangeFrom;

struct SquareMultiple {
    iter: RangeFrom<i64>,
    x: i64,
}

impl SquareMultiple {
    fn new(x: i64) -> Self {
        SquareMultiple {
            iter: (0..),
            x: x,
        }
    }
}

impl Iterator for SquareMultiple {
    type Item = i64;

    fn next(&mut self) -> Option<i64> {
        self.iter.next().map(|i| self.x * (self.x + i))
    }
}

fn main() {
    for i in SquareMultiple::new(5).take(10) {
        println!("{}", i);
    }
}