CODEWITHSUNDEEP
rust
The Unfun Cat
The Unfun Cat

Reputation: 32008

How do I create a Rust function that is generic in the operation applied?

I have a function called new_vec. It takes two vectors and creates a new one, by performing an elementwise operation on the pair of elements from the zipped vectors.

fn main() {
    let v1s = vec![1, 0, 1];
    let v2s = vec![0, 1, 1];
    let v3s = new_vec(v1s, v2s);
    println!("{:?}", v3s) // [1, 1, 2]
}

fn new_vec(v1s: Vec<i32>, v2s: Vec<i32>) -> Vec<i32> {
    let mut v3s = Vec::<i32>::new();
    for (v1, v2) in v1s.iter().zip(v2s.iter()) {
         v3s.push(v1 + v2) // would also like to use -
    }
    v3s
}

I want to have a new_vec function for the common binary operation that is possible to use on two integers, such as +, -, /, *.

How do I do this? I can imagine two ways: macros and closures. A minimal example of how to do this in the best way, for example with + and - would be appreciated.

Upvotes: 0

Views: 99

Answers (1)

ljedrz
ljedrz

Reputation: 22273

I would pass a closure:

fn new_vec<F>(v1s: &[i32], v2s: &[i32], foo: F) -> Vec<i32>
    where F: Fn(i32, i32) -> i32
{
    let mut v3s = Vec::<i32>::new();
    for (&v1, &v2) in v1s.iter().zip(v2s.iter()) {
        v3s.push(foo(v1, v2))
    }
    v3s
}

fn main() {
    let v1s = vec![1, 0, 1];
    let v2s = vec![0, 1, 1];
    let v3s = new_vec(&v1s, &v2s, |x, y| x - y);
    let v4s = new_vec(&v1s, &v2s, |x, y| x + y);
    println!("{:?}", v3s); // [1, -1, 0]
    println!("{:?}", v4s); // [1, 1, 2]
}

Note the change in the first two parameters; if your function doesn't need to consume its arguments, references are preferable to Vectors - in this case &[i32].

This implementation is not too efficient because the resulting Vector is extended incrementally; it's better if you modified it as follows to reduce the number of allocations:

fn new_vec<F>(v1s: &[i32], v2s: &[i32], foo: F) -> Vec<i32>
    where F: Fn(i32, i32) -> i32
{
    v1s.iter().zip(v2s.iter()).map(|(&x, &y)| foo(x, y)).collect()
}

Upvotes: 5

Related Questions