Use map object as function input

Question

I have a function that calculates the variance of an iterator of floats, I would like to be able to call this function on an iterator after using the map method to transform it.

use num::Float;

fn compute_var_iter<'a, I, T>(vals: I) -> T
where
    I: Iterator<Item = &'a T>,
    T: 'a + Float + std::ops::AddAssign,
{
    // online variance function
    // Var = E[X^2] - E[X]^2 
    // corrects for 1/n -> 1/(n-1)

    let mut x = T::zero();
    let mut xsquare = T::zero();
    let mut len = T::zero();

    for &val in vals {
        x += val;
        xsquare += val * val;
        len += T::one();
    }

    ((xsquare / len) - (x / len) * (x / len)) / (len - T::one()) * len
}

fn main() {
    let a: Vec<f64> = (1..100001).map(|i| i as f64).collect();
    let b: Vec<f64> = (0..100000).map(|i| i as f64).collect();

    dbg!(compute_var_iter(&mut a.iter())); // this works
    dbg!(compute_var_iter(a.iter().zip(b).map(|(x, y)| x * y))); // this does not work
}

Is there a performant way to get the map output back to an iterator or to make the function take the map object as an input so that we can avoid having to .collect() and keep the execution lazy?

Answer 1

You can use the iterator objects directly without collect:

use num::Float;

fn compute_var_iter<I, T>(vals: I) -> T
where
    I: Iterator<Item = T>,
    T: Float + std::ops::AddAssign,
{
    // online variance function
    // Var = E[X^2] - E[X]^2 
    // corrects for 1/n -> 1/(n-1)

    let mut x = T::zero();
    let mut xsquare = T::zero();
    let mut len = T::zero();

    for val in vals {
        x += val;
        xsquare += val * val;
        len += T::one();
    }

    ((xsquare / len) - (x / len) * (x / len)) / (len - T::one()) * len
}

fn main() {
    let a = (1..100001).map(|i| i as f64);
    let b = (0..100000).map(|i| i as f64);
    let c: Vec<f64> = (0..10000).map(|i| i as f64).collect();

    dbg!(compute_var_iter(a.clone())); // this works
    dbg!(compute_var_iter(c.iter().map(|i| *i))); // this works
    dbg!(compute_var_iter(a.zip(b).map(|(x, y)| x * y))); 
}

Playground

Notice that you would need to clone the iterator if you intend to use it several times. Also you do not really need to use references since numbers are usually Copy and the cost is the same as creating the references itself.