How do I tell the borrow checker this is safe? - image

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()
}