Is it possible to implement a custom Iterator::sum in Rust?

Question

I am trying to learn about iterators with an iterator that produces triangle numbers. Triangle numbers are 1, 3, 6, 10, 15 where 1 = 1, 3 = 1 + 2, 6 = 1 + 2 + 3 etc. I have the basics of this created:

pub struct Triangle {
    cur: u32,
    n: u32,
    m: u32,
}

impl Iterator for Triangle {
    type Item = u32;

    fn next(&mut self) -> Option<u32> {
        if self.n == self.m {
            return None;
        }
        self.n = self.n + 1;
        self.cur = self.cur + self.n;

        Some(self.cur)
    }
}

A quick runnable example of this is

let t = Triangle { cur: 0, n: 0, m: 10 };
let s: u32 = t.sum();
println!("{}", s);  // prints 220

Is it possible to create a custom summation function for the iterator that returns type u32? I was hoping to be able to do this with the default iterator and sum functions, and not have to make my own specialized function.

What I was hoping to be able to do is:

use std::iter::Sum;

impl Sum<u32> for u32 {
    fn sum<I>(iter: I) -> Self
    where
        I: Triangle,
    {
        let nsum = (self.n * (self.n + 1) * (self.n + 2)) / 6;
        let msum = (self.m * (self.m + 1) * (self.m + 2)) / 6;
        msum - nsum
    }
}

but this does not work. The error that I get with this is

error[E0404]: expected trait, found struct `Triangle`
  --> src/main.rs:26:12
   |
26 |         I: Triangle,
   |            ^^^^^^^^ not a trait

I could change it from Triangle to Iterator like it wants, but that would prevent me from accessing the m and n values of the Triangle struct.

How do I do this? Is it impossible? I know that I could write my own function like my_sum(), but I was hoping to be able to do it in the context of an iterator.

Answer 1

You cannot specialize the existing implementation of Sum, but you can specialize Iterator::sum in your iterator! It's a bit tricky, though, since its return type is generic.

use std::iter::{self, Sum};

impl Iterator for Triangle {
    // existing members are unchanged

    fn sum<S>(self) -> S
    where
        S: Sum<Self::Item>,
    {
        let nsum = (self.n * (self.n + 1) * (self.n + 2)) / 6;
        let msum = (self.m * (self.m + 1) * (self.m + 2)) / 6;
        S::sum(iter::once(msum - nsum))
    }
}

We cannot return a fixed type (such as u32), as that would not respect the contract defined by the Iterator trait. All we know about the return type S is that it implements Sum<Self::Item>. Sum has a single method, sum, that returns Self, so we can use it to produce a value of type S. The method expects an iterator; we feed it a Once, "an iterator that yields an element exactly once". Since that iterator will iterate a fixed number of times, we can expect sum to perform a fixed number of operations.

When you compile the program in release mode and S is u32, the whole sum call is optimized away, and the function returns msum - nsum directly.