How do I avoid using clone here?

Viewed 89
use rayon::prelude::*;
use image::RgbImage;


#[inline]
fn lerp(pct: f32, a: f32, b: f32) -> f32 {
    pct.mul_add(b - a, a)
}
#[inline]
fn distance(x: i32, y: i32) -> f32 {
    ((x * x + y * y) as f32).sqrt()
}

struct ColorCalculator {
    from: [f32; 3],
    to: [f32; 3],
    center_x: i32,
    center_y: i32,
    max_dist: f32,
}

impl ColorCalculator {
    fn new(from: [u8; 3], to: [u8; 3], width: u32, height: u32) -> Self {
        let center_x = width as i32 / 2;
        let center_y = height as i32 / 2;
        Self {
            from: from.map(|channel| channel as f32),
            to: to.map(|channel| channel as f32),
            center_x,
            center_y,
            max_dist: distance(center_x, center_y),
        }
    }
    fn calculate(&self, x: u32, y: u32) -> [u8; 3] {
        let x_dist = self.center_x - x as i32;
        let y_dist = self.center_y - y as i32;

        let t = distance(x_dist, y_dist) / self.max_dist;
    }
}

pub fn radial_gradient_mirror(
    geometry: [u32; 2],
    inner_color: [u8; 3],
    outer_color: [u8; 3],
) -> RgbImage {
    let [width, height] = geometry;

    let color_calculator = ColorCalculator::new(inner_color, outer_color, geometry[0], geometry[1]);

    let mut buf: Vec<_> = (0..(height / 2))
        .into_par_iter()
        .flat_map(|y| {
            let mut row: Vec<[u8; 3]> = Vec::with_capacity(width as usize);
            for x in 0..(width / 2) {
                row.push(color_calculator.calculate(x, y))
            }
            row.extend(row.clone().iter().rev());
            row
        })
        .collect();

    buf.extend(buf.clone().iter().rev());
    let buf = buf.into_iter().flatten().collect();

    RgbImage::from_raw(width, height, buf).unwrap()
}

I have an algorithm to generate a radial gradient. The original version would calculate the color of each pixel but because there is a horizontal and vertical line of symmetry I can calculate the colors for the top left corner and use some vector manipulation to mirror. I managed to do this by cloning and reversing the iterator buf.clone().iter().rev() although this is slow. How can I avoid this and are there any other optimizations I could use?

3 Answers

You do not really need to populate the buffer as later you override it, you can use direct chained iterators for doing so (copied is used to to have owned values), original buff will be droped at the end of the scope:

let i1 = buf.iter().copied();
let i2 = buf.iter().copied().rev();
let buf = i1.chain(i2).flatten().collect();

Full method:

pub fn radial_gradient_mirror(
    geometry: [u32; 2],
    inner_color: [u8; 3],
    outer_color: [u8; 3],
) -> RgbImage {
    let [width, height] = geometry;

    let color_calculator = ColorCalculator::new(inner_color, outer_color, geometry[0], geometry[1]);

    let mut buf: Vec<_> = (0..(height / 2))
        .into_par_iter()
        .flat_map(|y| {
            let mut row: Vec<[u8; 3]> = Vec::with_capacity(width as usize);
            for x in 0..(width / 2) {
                row.push(color_calculator.calculate(x, y))
            }
            row.extend(row.clone().iter().rev());
            row
        })
        .collect();

    let i1 = buf.iter().copied();
    let i2 = buf.iter().copied().rev();
    let buf = i1.chain(i2).flatten().collect();

    RgbImage::from_raw(width, height, buf).unwrap()
}

Playground

I believe your core question is how to optimize this:

buf.extend(buf.clone().iter().rev());

Vec has extend_from_within which can be used to extend a Vec from its own contents... however you want the results in reverse order. So I don't think there is a neat way using iterators to express this, since we're mutating and iterating over the same Vec at the same time. However doing this using a loop isn't too bad:

fn mirror_vec<T: Clone>(v: &mut Vec<T>) {
    let n = v.len();
    v.reserve(n);
    for i in (0..n).rev() {
        v.push(v[i].clone());
    }
}

Note that the clone is not the only performance bottleneck here, and probably not the main one either. A very big issue is using a Vec<Vec<_>> for a 2D array. It is much faster to use a single flat array and some indexing magic:

pub fn radial_gradient_mirror(
    geometry: [u32; 2],
    inner_color: [u8; 3],
    outer_color: [u8; 3],
) -> RgbImage {
    let width = geometry[0] as usize;
    let height = geometry[1] as usize;

    let color_calculator = ColorCalculator::new(inner_color, outer_color, geometry[0], geometry[1]);

    // Note: need to check that dimensions are even, or handle the case when they are odd! */
    let mut buf = Vec::with_capacity (width * height * 3);
    for r in 0..height/2 {
        for c in 0..width/2 {
            buf.extend_from_slice (&color_calculator.calculate (c, r));
        }
        for c in width/2+1 .. width {
            buf.extend_from_within ((r*width + width-c-1)*3 .. (r*width + width-c)*3);
        }
    }
    for r in height/2+1 .. height {
        buf.extend_from_within ((height-r-1) * width * 3 .. (height-r) * width * 3);
    }
    
    RgbImage::from_raw(width as u32, height as u32, buf).unwrap()
}

Playground

Note that as with all performance-related questions, you should benchmark in release mode to make sure that speed follows your expectations. In particular on modern computers, memory accesses are much slower than the CPU and it is often faster to re-compute a value rather than try to get it from a cached memory location!

Related