How to implement trait Add for custom type?

I have a type MyVec:

use std::ops::Add;

pub struct MyVec(Vec<f32>);

impl Add for MyVec {
    type Output = Map<_, _>;
    fn add(self, rhs: Self) -> Self::Output {
        self.0.iter()
            .zip(rhs.0.iter())
            .map(|(x, y)| x + y)
    }
}

I am trying to implement Add for MyVec, so my_vec1 + my_vec2 would return the type same as my_vec1.0.iter().zip(my_vec2.0.iter()).map(|(x, y)| x + y). But I failed to write down the type of Map.
I tried this type Output = Map<Zip<Iter<'a, f32>, Iter<'b, f32>>, fn((&f32, &f32)) -> f32>;, but it doesn't work.What should I do?

The troubles of writing down types involving closure types is a known issue, and the proposed and unstably implemented feature to address the issue is called type_alias_impl_trait. With it, you can decide on either the closure or the whole iterator as an abstraction point and write this code e.g. as

#![feature(impl_trait_in_assoc_type)]

use std::iter::{Map, Zip};
use std::ops::Add;
use std::slice::Iter;

pub struct MyVec(Vec<f32>);

impl<'a, 'b> Add<&'b MyVec> for &'a MyVec {
    type Output = Map<Zip<Iter<'a, f32>, Iter<'b, f32>>, impl FnMut((&'a f32, &'b f32)) -> f32>;
    fn add(self, rhs: &'b MyVec) -> Self::Output {
        self.0.iter().zip(rhs.0.iter()).map(|(x, y)| x + y)
    }
}

Note that the feature seems to have already been split up into parts, reading into the tracking issue, it seems like the impl_trait_in_assoc_type-part is supposed to stabilize more quickly.

Also, you can see that I've changed the type signatures so that self and rhs are references. The returned iterator does hold borrows to the Vecs, and thus your code as-is would run into use-after-free issues prevented by the borrow checked.

The alternative of using something like type Output = impl Iterator<Item = f32>; is also a possibility, though there’s some issues around the ergonomics of capturing lifetimes in such a type that I don't want to discuss in detail.


If your goal is to keep the arguments owned instances of MyVec, then I don't see much of a need to use an iterator as the return type. If you do still want it, it's going to have you use into_iter() instead of iter() to avoid the borrowing issue.

#![feature(impl_trait_in_assoc_type)]

use std::ops::Add;

pub struct MyVec(Vec<f32>);

impl Add for MyVec {
    type Output = impl Iterator<Item = f32>;
    fn add(self, rhs: MyVec) -> Self::Output {
        self.0.into_iter().zip(rhs.0).map(|(x, y)| x + y)
    }
}

The only stable way of returning an abstract iterator type right now is to use a boxed trait object, which would unfortunately come with minor performance implications due to boxing and dynamic dispatch, but it would look as follows:

use std::ops::Add;

pub struct MyVec(Vec<f32>);

impl Add for MyVec {
    type Output = Box<dyn Iterator<Item = f32>>;
    fn add(self, rhs: MyVec) -> Self::Output {
        Box::new(self.0.into_iter().zip(rhs.0).map(|(x, y)| x + y))
    }
}

It's also possible to introduce dynamic dispatch only into the mapped function, e.g. in this case via a function pointer type

use std::iter::{Map, Zip};
use std::ops::Add;
use std::vec::IntoIter; 

pub struct MyVec(Vec<f32>);

impl Add for MyVec {
    type Output = Map<Zip<IntoIter<f32>, IntoIter<f32>>, fn((f32, f32)) -> f32>;
    fn add(self, rhs: MyVec) -> Self::Output {
        self.0.into_iter().zip(rhs.0).map(|(x, y)| x + y)
    }
}

And defining a custom iterator type would also always be an option, replacing (and at least partially re-implementing) the Map type for the specific case of addition of f32. (No code example here.)


The simplest approach might be to simply return a MyVec again though. This does all of the addition in advance, but I don't see the huge appeal of using an iterator instead, especially if you're working with owned types anyways.

use std::ops::Add;

pub struct MyVec(Vec<f32>);

impl Add for MyVec {
    type Output = MyVec;
    fn add(self, rhs: MyVec) -> Self::Output {
        MyVec(self.0.into_iter().zip(rhs.0).map(|(x, y)| x + y).collect())
    }
}

It's quite possible that some specialization magic in the compiler already optimizes away the re-allocation in this code. A manually in-place version would look e.g. as follows:

use std::ops::Add;

pub struct MyVec(Vec<f32>);

impl Add for MyVec {
    type Output = MyVec;
    fn add(mut self, rhs: MyVec) -> Self::Output {
        self.0.truncate(rhs.0.len());
        self.0.iter_mut().zip(rhs.0).for_each(|(x, y)| *x += y);
        self
    }
}

either way, you can also consider taking only self owned and keeping rhs by-reference, in this in-place approach (with either of the last two code examples above), the rhs is not needed to be owned

use std::ops::Add;

pub struct MyVec(Vec<f32>);

impl Add<&MyVec> for MyVec {
    type Output = MyVec;
    fn add(mut self, rhs: &MyVec) -> Self::Output {
        self.0.truncate(rhs.0.len());
        self.0.iter_mut().zip(&rhs.0).for_each(|(x, y)| *x += y);
        self
    }
}

In this case, you could also implement += and then do + in terms of it:


pub struct MyVec(Vec<f32>);

impl AddAssign<&MyVec> for MyVec {
    fn add_assign(&mut self, rhs: &MyVec) {
        self.0.truncate(rhs.0.len());
        self.0.iter_mut().zip(&rhs.0).for_each(|(x, y)| *x += y);
    }
}

impl Add<&MyVec> for MyVec {
    type Output = MyVec;
    fn add(mut self, rhs: &MyVec) -> Self::Output {
        self += rhs;
        self
    }
}
1 Like