2024-01-27

Escaping Closure Ownership Conundrums in Rust

I’ve noticed a common issue people bump into (including myself) more than once now when starting out with Rust, and thought I’d share what I’ve come to see as the natural solution!

For illustration purposes, let’s say you have this function that simulates a growing garden of plants:

fn grow_garden_for_a_year(mut garden: Vec<u8>) {
  for plant_height in &mut garden {
    *plant_height += 4; //summer growth
  }
  for plant_height in &mut garden {
    *plant_height += 2; //fall growth
  }
  for plant_height in &mut garden {
    *plant_height += 1; //winter growth
  }
  for plant_height in &mut garden {
    *plant_height += 3; //spring growth
  }

  println!("at the end of the year, plant heights are {:?}", garden);
}

one natural simplification is to make a small closure that can grow all plants some specified amount:

fn grow_garden_for_a_year(mut garden: Vec<u8>) {
  let mut grow_all = |rate| {
    for plant_height in &mut garden {
      *plant_height += rate;
    }
  };
  grow_all(4); //summer growth
  grow_all(2); //fall growth
  grow_all(1); //winter growth
  grow_all(3); //spring growth

  println!("at the end of the year, plant heights are {:?}", garden);
}

Let’s say you now want to build out the model to simulate hurricanes, that sometimes destroy a plant between fall and winter, and you try to do this:

fn grow_garden_for_a_year(mut garden: Vec<u8>) {
  let mut grow_all = |rate| {
    for plant_height in &mut garden {
      *plant_height += rate;
    }
  };
  grow_all(4); //summer growth
  grow_all(2); //fall growth
  garden[0] = 0; //hurricane destroyed least protected plant!
  grow_all(1); //winter growth
  grow_all(3); //spring growth

  println!("at the end of the year, plant heights are {:?}", garden);
}

Then rust blows up in your face!

error[E0499]: cannot borrow `garden` as mutable more than once at a time
  --> src/lib.rs:11:3
   |
4  |   let mut grow_all = |rate| {
   |                      ------ first mutable borrow occurs here
5  |     for plant_height in &mut garden {
   |                              ------ first borrow occurs due to use of `garden` in closure
...
11 |   garden[0] = 0; //hurricane destroyed least protected plant!
   |   ^^^^^^ second mutable borrow occurs here
12 |   grow_all(1); //winter growth
   |   -------- first borrow later used here

For more information about this error, try `rustc --explain E0499`.

The closure we created does use a mutable reference to garden, and keeps it until the function isn’t called anymore. Trying to use the variable somewhere else before this closure isn’t used anymore would violate rusts rules around only allowing one mutable borrow. But what should we do in these kinds of situations then?

One approach I’ve seen is to just make it a function that takes a mutable borrow to the garden, this way, it can give back the variable between calls:

fn grow_garden_for_a_year(mut garden: Vec<u8>) {
  fn grow_all(garden: &mut [u8], rate: u8) {
    for plant_height in garden {
      *plant_height += rate;
    }
  }
  grow_all(&mut garden, 4); //summer growth
  grow_all(&mut garden, 2); //fall growth
  garden[0] = 0; //hurricane destroyed least protected plant!
  grow_all(&mut garden, 1); //winter growth
  grow_all(&mut garden, 3); //spring growth

  println!("at the end of the year, plant heights are {:?}", garden);
}

I’d say this is okay, but feels a bit clumsy and isn’t great if you’re modifying more than one variable. Another approach I’ve seen is to create a macro:

fn grow_garden_for_a_year(mut garden: Vec<u8>) {

    macro_rules! grow_all {
        ($rate:expr) => {
            for plant_height in &mut garden {
                *plant_height += $rate;
            }
        };
    }

    grow_all!(4); // Summer growth
    grow_all!(2); // Fall growth
    garden[0] = 0; // Hurricane destroyed least protected plant!
    grow_all!(1); // Winter growth
    grow_all!(3); // Spring growth

    println!("At the end of the year, plant heights are {:?}", garden);
}

This doesn’t feel great either to me. I expect macros to perform magic when needed, but shoveling dirt with them feels wrong (and potentially creates a lot of code duplication, even if it’s hidden in this case).

More often than not, it might be a sign that you want some behaviors associated with some collection of data, exactly what structs are for!

fn grow_garden_for_a_year(mut garden: Vec<u8>) {
    
    struct Garden { heights: Vec<u8> }
    
    impl Garden {
        fn grow_all(&mut self, rate: u8) {
            for plant_height in &mut self.heights {
                *plant_height += rate;
            }
        }
    }
    
    let mut g = Garden { heights: garden };

    g.grow_all(4); // Summer growth
    g.grow_all(2); // Fall growth
    g.heights[0] = 0; // Hurricane destroyed least protected plant!
    g.grow_all(1); // Winter growth
    g.grow_all(3); // Spring growth

    println!("At the end of the year, plant heights are {:?}", g.heights);
}

If this is a construct useful outside of this function, you might want to do this:

struct Garden { heights: Vec<u8> }
    
impl Garden {
    fn grow_all(&mut self, rate: u8) {
        for plant_height in &mut self.heights {
            *plant_height += rate;
        }
    }
    
    fn hurricane(&mut self) {
        self.heights[0] = 0;
    }
}
    
fn grow_garden_for_a_year(mut garden: Garden) {

    garden.grow_all(4); // Summer growth
    garden.grow_all(2); // Fall growth
    garden.hurricane(); // Hurricane destroyed least protected plant!
    garden.grow_all(1); // Winter growth
    garden.grow_all(3); // Spring growth

    println!("At the end of the year, plant heights are {:?}", garden.heights);
}

Even if it is a bit more verbose, I think it’s a great thing that Rust almost forces the realization that you are dealing with a collection of data that has some associated behaviors, and that this can be better expressed.

Let me know your thoughts!