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

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.