How do I tell the borrow checker this is safe?

Question

use image::{Rgb, RgbImage};
use rayon::prelude::*;

#[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;

        [
            lerp(t, self.from[0], self.to[0]) as u8,
            lerp(t, self.from[1], self.to[1]) as u8,
            lerp(t, self.from[2], self.to[2]) as u8,
        ]
    }
}

fn radial_gradient(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, width, height);

    let mut background = RgbImage::new(width, height);

    (0..height / 2).into_par_iter().for_each(|y| {
        for x in 0..width / 2 {
            let color = Rgb(color_calculator.calculate(x, y));
            background.put_pixel(x, y, color);
            background.put_pixel(width - x - 1, y, color);
            background.put_pixel(x, height - y - 1, color);
            background.put_pixel(width - x - 1, height - y - 1, color);
        };
    });
    
    background
}

I know that I could just use a mutex here although it is unnecessary since provided my code is correct no pixel should be mutated more than once. So how do I tell rust that doing background.put_pixel(x, y, color) in multiple threads is actually okay here?

I'm guessing some use of unsafe has to be used here although I am new to rust and am not sure how to use it effectively here.

Here's the error

error[E0596]: cannot borrow `background` as mutable, as it is a captured variable in a `Fn` closure
   --> src\lib.rs:212:13
    |
212 |             background.put_pixel(x, y, color);
    |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ cannot borrow as mutable

error[E0596]: cannot borrow `background` as mutable, as it is a captured variable in a `Fn` closure
   --> src\lib.rs:213:13
    |
213 |             background.put_pixel(width - x - 1, y, color);
    |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ cannot borrow as mutable

error[E0596]: cannot borrow `background` as mutable, as it is a captured variable in a `Fn` closure
   --> src\lib.rs:214:13
    |
214 |             background.put_pixel(x, height - y - 1, color);
    |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ cannot borrow as mutable

error[E0596]: cannot borrow `background` as mutable, as it is a captured variable in a `Fn` closure
   --> src\lib.rs:215:13
    |
215 |             background.put_pixel(width - x - 1, height - y - 1, color);
    |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ cannot borrow as mutable

Answer 1

You can't. At least not with an RgbImage .

put_pixel takes a &mut self . In Rust, it's undefined behavior to have two &mut references alias - the optimizer can do some funky stuff to your code if you break this assumption.

You will probably have an easier time creating a Vec<u8> of pixel data, calculating each pixel's value using Rayon's parallel iterators (which will take special care to not alias the mutable references), then assemble the buffer into an image using from_vec .

Answer 2

You can't do this with RgbImage (or any ImageBuffer ), but you can do it if you work on a raw Vec<u8> in pure safe code.

Essentially the idea is to use split_at_mut and par_(r)chunks_exact_mut to produce parallel iterators that start from each corner of the image.

First, we allocate a chunk of memory

fn radial_gradient(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, width, height);
    
    // assertions here to hopefully help the optimizer later
    assert!(width % 2 == 0);
    assert!(height % 2 == 0);

    // allocate memory for the image
    let mut background = RgbImage::new(width, height).into_vec();

Then, we split the memory into a top and bottom half, and create iterators over each row, starting at the top and bottom

    let width = width as usize;
    let height = height as usize;

    // split background into top and bottom, 
    // so we can utilize the x-axis symmetry to reduce color calculations
    let (top, bottom) = background.split_at_mut(width * height * 3 / 2);

    // use chunks to split each half into rows,
    // so we can utilize the y-axis symmetry to reduce color calculations
    let top_rows = top.par_chunks_exact_mut(width * 3);
    let bottom_rows = bottom.par_rchunks_exact_mut(width * 3);

Then, we zip those iterators together, so we iterate over the top row and bottom row together, then the next row in on each side, etc. Add enumerate to get the Y coordinate and then we will flat_map so our pixel iterator later gets unwrapped into the main iterator.

    // zip to iterate over top and bottom row together
    // enumerate to get the Y coordinate
    top_rows.zip(bottom_rows).enumerate().flat_map(|(y, (top_row, bottom_row))| {

Then, split each row at the middle so we can have four iterators, one for each corner

        // split each row at the y-axis
        let (tl, tr) = top_row.split_at_mut(width * 3 / 2);
        let (bl, br) = bottom_row.split_at_mut(width * 3 / 2);
        
        // iterate over pixels (chunks of 3 bytes) from the
        // top left, bottom left, top right, and bottom right half-rows
        let tl = tl.par_chunks_exact_mut(3);
        let bl = bl.par_chunks_exact_mut(3);
        let tr = tr.par_rchunks_exact_mut(3);
        let br = br.par_rchunks_exact_mut(3);

Then, zip the four pixel iterators together, so we iterate from the four corners simultaneously

        // zip to iterate over each set of four pixels together
        // enumerate to get the X coordinate
        tl.zip_eq(bl).zip_eq(
            tr.zip_eq(br)
        ).enumerate().map(move |(x, ((tl, bl), (tr, br)))| {
            // add the y coordinate to the pixel-wise iterator
            ((x, y), (tl, bl, tr, br))
        })

Then, iterate over each set of four pixels, copying the color into each And convert back into an RgbImage

    }).for_each(|((x, y), (tl, bl, tr, br))| {
        // copy the color into the four symmetric pixels
        let color = color_calculator.calculate(x as u32, y as u32);
        tl.copy_from_slice(&color);
        bl.copy_from_slice(&color);
        tr.copy_from_slice(&color);
        br.copy_from_slice(&color);
    });

    RgbImage::from_vec(width as u32, height as u32, background).unwrap()
}

playground

It's hard to say if this will be more or less performant than the strategy of coloring one quadrant and copying it to the rest, but it's worth a try. It may also not be worth the cognitive overhead of all of the chunking wizardry going on.

---

Full code

fn radial_gradient(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, width, height);
    
    // assertions here to hopefully help the optimizer later
    assert!(width % 2 == 0);
    assert!(height % 2 == 0);

    // allocate memory for the image
    let mut background = RgbImage::new(width, height).into_vec();
    
    let width = width as usize;
    let height = height as usize;
    
    // split background into top and bottom, 
    // so we can utilize the x-axis symmetry to reduce color calculations
    let (top, bottom) = background.split_at_mut(width * height * 3 / 2);
    
    // use chunks to split each half into rows,
    // so we can utilize the y-axis symmetry to reduce color calculations
    let top_rows = top.par_chunks_exact_mut(width * 3);
    let bottom_rows = bottom.par_rchunks_exact_mut(width * 3);
    
    // zip to iterate over top and bottom row together
    // enumerate to get the Y coordinate
    top_rows.zip(bottom_rows).enumerate().flat_map(|(y, (top_row, bottom_row))| {
        // split each row at the y-axis
        let (tl, tr) = top_row.split_at_mut(width * 3 / 2);
        let (bl, br) = bottom_row.split_at_mut(width * 3 / 2);
        
        // iterate over pixels (chunks of 3 bytes) from the
        // top left, bottom left, top right, and bottom right half-rows
        let tl = tl.par_chunks_exact_mut(3);
        let bl = bl.par_chunks_exact_mut(3);
        let tr = tr.par_rchunks_exact_mut(3);
        let br = br.par_rchunks_exact_mut(3);
        
        // zip to iterate over each set of four pixels together
        // enumerate to get the X coordinate
        tl.zip_eq(bl).zip_eq(
            tr.zip_eq(br)
        ).enumerate().map(move |(x, ((tl, bl), (tr, br)))| {
            // add the y coordinate to the pixel-wise iterator
            ((x, y), (tl, bl, tr, br))
        })
    }).for_each(|((x, y), (tl, bl, tr, br))| {
        // copy the color into the four symmetric pixels
        let color = color_calculator.calculate(x as u32, y as u32);
        tl.copy_from_slice(&color);
        bl.copy_from_slice(&color);
        tr.copy_from_slice(&color);
        br.copy_from_slice(&color);
    });
    
    RgbImage::from_vec(width as u32, height as u32, background).unwrap()
}

Answer 3

You can use a raw pointer, cast it to unsigned integer, and then unsafely dereference inside the loop:

let background: *mut RgbImage =
    &mut RgbImage::new(width,height) as *mut RgbImage;
let pointer = background as usize;
(0..height / 2).into_par_iter().for_each(|y| {
    for x in 0..width / 2 {
        ...
        let mut background = unsafe { (pointer as *mut RgbImage).as_mut() }.unwrap();
        ...
    };
});

unsafe { background.as_ref().unwrap() })