multithreadingrustlifetimeinterior-mutability

Rust: allow multiple threads to modify an image (wrapper of a vector)?


Suppose I have a "image" struct that wraps a vector:

type Color = [f64; 3];

pub struct RawImage
{
    data: Vec<Color>,
    width: u32,
    height: u32,
}

impl RawImage
{
    pub fn new(width: u32, height: u32) -> Self
    {
        Self {
            data: vec![[0.0, 0.0, 0.0]; (width * height) as usize],
            width: width,
            height: height
        }
    }

    fn xy2index(&self, x: u32, y: u32) -> usize
    {
        (y * self.width + x) as usize
    }
}

It is accessible through a "view" struct, which abstracts an inner block of the image. Let's assume that I only want to write to the image (set_pixel()).

pub struct RawImageView<'a>
{
    img: &'a mut RawImage,
    offset_x: u32,
    offset_y: u32,
    width: u32,
    height: u32,
}

impl<'a> RawImageView<'a>
{
    pub fn new(img: &'a mut RawImage, x0: u32, y0: u32, width: u32, height: u32) -> Self
    {
        Self{ img: img,
              offset_x: x0, offset_y: y0,
              width: width, height: height, }
    }

    pub fn set_pixel(&mut self, x: u32, y: u32, color: Color)
    {
        let index = self.img.xy2index(x + self.offset_x, y + self.offset_y);
        self.img.data[index] = color;
    }
}

Now suppose I have an image, and I want to have 2 threads modifying it at the same time. Here I use rayon's scoped thread pool:

fn modify(img: &mut RawImageView)
{
    // Do some heavy calculation and write to the image.
    img.set_pixel(0, 0, [0.1, 0.2, 0.3]);
}

fn main()
{
    let mut img = RawImage::new(20, 10);
    let pool = rayon::ThreadPoolBuilder::new().num_threads(2).build().unwrap();
    pool.scope(|s| {
        let mut v1 = RawImageView::new(&mut img, 0, 0, 10, 10);
        let mut v2 = RawImageView::new(&mut img, 10, 0, 10, 10);
        s.spawn(|_| {
            modify(&mut v1);
        });
        s.spawn(|_| {
            modify(&mut v2);
        });
    });
}

This doesn't work, because

  1. I have 2 &mut img at the same time, which is not allowed
  2. "closure may outlive the current function, but it borrows v1, which is owned by the current function"

So my questions are

  1. How can I modify RawImageView, so that I can have 2 threads modifying my image?
  2. Why does it still complain about life time of the closure, even though the threads are scoped? And how do I overcome that?

Playground link

One approach that I tried (and it worked) was to have modify() just create and return a RawImage, and let the thread push it into a vector. After all the threads were done, I constructed the full image from that vector. I'm trying to avoid this approach due to its RAM usage.


Solution

  • Your two questions are actually unrelated.

    First the #2 that is easier:

    The idea of the Rayon scoped threads is that the threads created inside cannot outlive the scope, so any variable created outside the scope can be safely borrowed and its references sent into the threads. But your variables are created inside the scope, and that buys you nothing.

    The solution is easy: move the variables out of the scope:

        let mut v1 = RawImageView::new(&mut img, 0, 0, 10, 10);
        let mut v2 = RawImageView::new(&mut img, 10, 0, 10, 10);
        pool.scope(|s| {
            s.spawn(|_| {
                modify(&mut v1);
            });
            s.spawn(|_| {
                modify(&mut v2);
            });
        });
    

    The #1 is trickier, and you have to go unsafe (or find a crate that does it for you but I found none). My idea is to store a raw pointer instead of a vector and then use std::ptr::write to write the pixels. If you do it carefully and add your own bounds checks it should be perfectly safe.

    I'll add an additional level of indirection, probably you could do it with just two but this will keep more of your original code.

    The RawImage could be something like:

    pub struct RawImage<'a>
    {
        _pd: PhantomData<&'a mut Color>,
        data: *mut Color,
        width: u32,
        height: u32,
    }
    impl<'a> RawImage<'a>
    {
        pub fn new(data: &'a mut [Color], width: u32, height: u32) -> Self
        {
            Self {
                _pd: PhantomData,
                data: data.as_mut_ptr(),
                width: width,
                height: height
            }
        }
    }
    

    And then build the image keeping the pixels outside:

        let mut pixels = vec![[0.0, 0.0, 0.0]; (20 * 10) as usize];
        let mut img = RawImage::new(&mut pixels, 20, 10);
    

    Now the RawImageView can keep a non-mutable reference to the RawImage:

    pub struct RawImageView<'a>
    {
        img: &'a RawImage<'a>,
        offset_x: u32,
        offset_y: u32,
        width: u32,
        height: u32,
    }
    

    And use ptr::write to write the pixels:

        pub fn set_pixel(&mut self, x: u32, y: u32, color: Color)
        {
            let index = self.img.xy2index(x + self.offset_x, y + self.offset_y);
            //TODO! missing check bounds
            unsafe { self.img.data.add(index).write(color) };
        }
    

    But do not forget to either do check bounds here or mark this function as unsafe, sending the responsibility to the user.

    Naturally, since your function keeps a reference to a mutable pointer, it cannot be send between threads. But we know better:

    unsafe impl Send for RawImageView<'_> {}
    

    And that's it! Playground. I think this solution is memory-safe, as long as you add code to enforce that your views do not overlap and that you do not go out of bounds of each view.