Make a new iterator trait to replace &[D] or (&[f64], &[f64])

Hi, hoping the community can help.

I'd like to consolidate some functions that only differ in their arguments. Here, I've made a trait (Direction) that can abstract across a custom type (Cartestian) and (f64, f64).

#[derive(Clone, Copy)]
struct Cartesian {
    x: f64,
    y: f64,
}
impl Cartesian {
    fn new(x: f64, y: f64) -> Self {
        Self { x, y }
    }
}

trait Direction: Copy {
    fn get_x(self) -> f64;
    fn get_y(self) -> f64;
}

impl Direction for Cartesian {
    fn get_x(self) -> f64 {
        self.x
    }

    fn get_y(self) -> f64 {
        self.y
    }
}

impl Direction for (f64, f64) {
    fn get_x(self) -> f64 {
        self.0
    }

    fn get_y(self) -> f64 {
        self.1
    }
}

struct Beam;
impl Beam {
    fn calc<D: Direction>(&self, dir: D) -> f64 {
        // Contrived for example's benefit
        dir.get_x() + dir.get_y()
    }
}

fn main() {
    let b = Beam;
    println!("{}", b.calc(Cartesian::new(5.0, 8.0)));
    println!("{}", b.calc((5.0, 8.0)));
}

This works well.

However my problem lies with a "calc many" function for Beam; it is allowed to take a slice of Cartesian or two slices of f64 to get multiple coordinate pairs. (The actual implementation iterates over the input coords, so it seems best to make an iterator trait but I may be wrong.)

Here is an attempt:

impl Beam {
    // This is in addition to the other function
    fn calc_many<D: Direction, T>(&self, dirs: &T) -> Vec<f64>
    where
        for<'a> &'a T: IntoIterator<Item = &'a D>,
    {
        dirs.into_iter()
            .map(|dir| dir.get_x() + dir.get_y())
            .collect()
    }
}

fn main() {
    let v = vec![
        Cartesian::new(1.0, 1.0),
        Cartesian::new(1.0, 2.0),
        Cartesian::new(2.0, 3.0),
    ];
    println!("{:?}", b.calc_many(&v));

    let a = [
        Cartesian::new(1.0, 1.0),
        Cartesian::new(1.0, 2.0),
        Cartesian::new(2.0, 3.0),
    ];
    println!("{:?}", b.calc_many(&a));

    let a = [(1.0, 1.0), (1.0, 2.0), (2.0, 3.0)];
    println!("{:?}", b.calc_many(&a));

    let a = [1.0, 1.0, 2.0];
    let a2 = [1.0, 2.0, 3.0];
    let i = a
        .into_iter()
        .zip(a2.into_iter())
        .map(|(a, b)| Cartesian::new(a, b));
    println!("{:?}", b.calc_many(i));
}

This almost works, but attempting to use i at the bottom of main doesn't work. Besides, I'd like users to be able to pass in slices instead of iterators.

What's the best way to make a "calc many" function that can take a slice of an appropriate input (here Direction), or a tuple of slices (e.g. (&[1.0, 1.0, 2.0], &[1.0, 2.0, 3.0])? I have tried to make another trait without much luck.

(For additional context, here's the code I'm trying to consolidate. I want to make any new code able to be used with rayon but I figured a normal iterator will help me understand how this works.)

You put IntoIterator as a bound. All you need to make your code work is to collect into a collection that implements it, like a Vector:

let i = a
        .into_iter()
        .zip(a2.into_iter())
        .map(|(a, b)| Cartesian::new(a, b)).collect::<Vec<_>>();
    println!("{:?}", b.calc_many(&i));

I would question that design decision for ergonomic reasons. An array of slices instead of a tuple would be a better option, since you can iterate through it.

Thanks for the reply. I'd like to avoid collecting if possible. Also, I could live with an array of slices over a tuple of slices, but is the tuple not possible?

The last line doesn't compile because the bound IntoIterator<Item = &'a D> requires that &D is owned by something other than the iterator (like a vector, as already pointed out). You probably don't want to create a vector just for iterating over the values.

One approach is to make calc_many even more generic, so that it accepts anything that can be borrowed as Direction. You can now call calc_many with owning iterators and non-owning iterators, even slices.

fn calc_many<T, I, D>(&self, dirs: T) -> Vec<f64>
where
    T: IntoIterator<Item = I>,
    I: Borrow<D>,
    D: Direction,
{
    dirs.into_iter()
        .map(|dir| dir.borrow().get_x() + dir.borrow().get_y())
        .collect()
}

The downside is that type annotations are required: Playground

2 Likes

Thanks @danfabo! This is much better than what I could come up with (as well as LLMs...)

As this level of generics is still a bit advanced for me, I'm curious what other solutions are out there. But for now I will accept this as the solution, thanks again!

One solution is to add a new blanket impl of Direction to make it work for all references:

impl<D: Direction + ?Sized> Direction for &D {
    fn get_x(self) -> f64 {
        (*self).get_x()
    }

    fn get_y(self) -> f64 {
        (*self).get_y()
    }
}

This lets us use a basic IntoIterator for calc_many(), without needing any type annotations when we call it (Rust Playground):

impl Beam {
    fn calc_many<D, I>(&self, dirs: I) -> Vec<f64>
    where
        D: Direction,
        I: IntoIterator<Item = D>,
    {
        dirs.into_iter()
            .map(|dir| dir.get_x() + dir.get_y())
            .collect()
    }
}

It's a bit trickier to allow tuples of slices directly. Already, we can call b.calc_many(iter::zip(tuple.0, tuple.1)), but to make it work without the extra zip() call would require something like a custom Directions trait that can turn either a slice or tuple into an Iterator.

1 Like

Just wanted to say thanks, this has helped my solution even more!