How to implement the extension trait pattern for functions that return `impl Trait`?

Question

Suppose I have the following helper method that adds two iterators (or IntoIterators) together using vector addition:

fn add<I1, I2, T>(a: I1, b: I2) -> impl Iterator<Item = T>
where
    I1: IntoIterator<Item = T>,
    I2: IntoIterator<Item = T>,
    T: std::ops::Add<Output = T>,
{
    a.into_iter().zip(b).map(|(x, y)| x + y)
}

It can be used as follows:

fn main() {
    let a = vec![1, 2, 3, 4, 5]; // Can use vectors
    let b = (6..=10).map(|x| x + 1); // Or other complicated iterators
    let sum_iter = add(a, b); // Can be used as an iterator if you wish to chain other functions
    let sum_vec = sum_iter.collect::<Vec<i32>>(); // Or you can collect it
    assert_eq!(sum_vec, vec![8, 10, 12, 14, 16]);
}

After implementing this, I decided I should change the API so that I can use the syntax a.add(b) which would make chaining these functions really nice. I attempted to implement the extension trait pattern but quickly realized this wouldn't be possible since impl Trait is not supported for return types of trait functions, and my add function is unable to specify a concrete type due to the closure |(x, y)| x + y that is passed into .map()

To help clarify my idea, this is essentially what I tried (obviously this doesn't compile):

trait IntoIteratorMathExt
where
    Self: IntoIterator,
{
    fn add<I>(self, other: I) -> impl Iterator<Item = Self::Item> 
    where
        I: IntoIterator<Item = Self::Item>,
        Self::Item: std::ops::Add<Output = Self::Item>;
}

impl<T: IntoIterator> IteratorMathExt for T {
    fn add<I>(self, other: I) -> impl Iterator<Item = Self::Item>
    where
        I: IntoIterator<Item = Self::Item>,
        Self::Item: std::ops::Add<Output = Self::Item>,
    {
        self.into_iter().zip(other).map(|(x, y)| x + y)
    }
}

fn main() {
    let a = vec![1, 2, 3, 4, 5];
    let b = (6..=10).map(|x| x + 1);
    let sum_iter = a.add(b); // <- Really nice syntax!
    let sum_vec = sum_iter.collect::<Vec<i32>>();
    assert_eq!(sum_vec, vec![8, 10, 12, 14, 16]);
}

Unsurprisingly, this gives the following compiler error:

error[E0562]: `impl Trait` not allowed outside of function and inherent method return types
  --> src\main.rs:34:34
   |
34 |     fn add<I>(self, other: I) -> impl Iterator<Item = Self::Item>
   |                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

error[E0562]: `impl Trait` not allowed outside of function and inherent method return types
  --> src\main.rs:41:34
   |
41 |     fn add<I>(self, other: I) -> impl Iterator<Item = Self::Item>
   |                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

I understand why impl Trait is not supported in trait methods, however, my trait is "practically" only ever implemented in one place so I'm hoping a workaround may exist.

So giving up on this approach, I'd like to know what the "correct" way to implement this pattern is.

The goal is to be able to use iterators a and b (either Iterator or IntoIterator, I'm not 100% what is the best), then be able to write a.add(b) and have it return an Iterator, with the same associated Item type as both a and b.

Any suggestions?

Aplet123 Aplet123 · Accepted Answer · 2021-06-01T01:22:11

Since your closure doesn't capture any values, it can actually be expressed as a function pointer type, which allows you to express the type (it's ugly, but it works):

type MathExtType<A, B, T> = std::iter::Map<std::iter::Zip<A, B>, fn((T, T)) -> T>;

trait IntoIteratorMathExt
where
    Self: IntoIterator,
{
    fn add<I>(self, other: I) -> MathExtType<Self::IntoIter, I::IntoIter, Self::Item>
    where
        I: IntoIterator<Item = Self::Item>,
        Self::Item: std::ops::Add<Output = Self::Item>;
}

impl<T: IntoIterator> IntoIteratorMathExt for T {
    fn add<I>(self, other: I) -> MathExtType<Self::IntoIter, I::IntoIter, Self::Item>
    where
        I: IntoIterator<Item = Self::Item>,
        Self::Item: std::ops::Add<Output = Self::Item>,
    {
        self.into_iter().zip(other).map(|(x, y)| x + y)
    }
}

fn main() {
    let a = vec![1, 2, 3, 4, 5];
    let b = (6..=10).map(|x| x + 1);
    let sum_iter = a.add(b); // <- Really nice syntax!
    let sum_vec = sum_iter.collect::<Vec<i32>>();
    assert_eq!(sum_vec, vec![8, 10, 12, 14, 16]);
}

How to implement the extension trait pattern for functions that return `impl Trait`?

2 Answers