Rust – Fast manipulation of a vector behind a HashMap using RefCell

Let’s analyze a function which maintains some 3×1 f64 vectors for plotting. There are multiple vectors and each of them is behind an enum key of a HashMap. Hasmap is a member property of a struct, so it looks something like this:

#[derive(Debug, Clone)]
pub struct Particle {
    pub pos: Vector3<f64>, // current Euclidean XYZ position [m]
    pub v: Vector3<f64>,   // current Euclidean XYZ accelleration [m/s]
    pub a: Vector3<f64>,   // current Euclidean XYZ accelleration [m/s^2]
    pub m: f64,            // particle mass [kg]
    pub r: f64,            // radius of the particle [m]
    pub e: f64,            // potential energy [kcal/mol]
    pub q: f64,            // electrical charge [C]
    pub particle_type: SubatomicParticleType,

    data_trace: HashMap<ParticleProperty, Vec<Matrix3x1<f64>>>
}

We have a function log_debug_data which will record the current position of a particle into the data_trace HashMap for the key ParticleProperty::Position, and it looks like this:

 pub fn log_debug_data(&mut self) {
        match self.data_trace.get(&ParticleProperty::Position) {
            Some(trace) => { 
                let mut t = trace.clone();
                t.push(self.pos);
                self.data_trace.insert(ParticleProperty::Position, t);
            },
            None => { self.data_trace.insert(ParticleProperty::Position, vec![self.pos]); }
        }
    }

We use pattern matching to check if the key already exists in the map so we know if a new vector should be created with a single entry or an element should be appended to an existing vector.

As you can see, appending the element to an existing vector had to be done through cloning the vector first, to obtain the mutable reference to it. If we tried to get the mutable reference without the cloning, Rust would complain about data_trace being of type Vec<>, which doesn’t implement the Copy trait needed for generating mutable reference:

Even if somehow we managed to make the compiler happy, the fact we got the mutable reference of something immutable must have meant Rust copied it under the hood, which is something we were aiming to avoid in the first place. So, what can we do?

Internal mutability to the rescue!

To avoid copying, we must change out struct so the data_trace property’s vectors are each wrapped with a RefCell<>, like this:

#[derive(Debug, Clone)]
pub struct Particle {
    pub pos: Vector3<f64>, // current Euclidean XYZ position [m]
    pub v: Vector3<f64>,   // current Euclidean XYZ accelleration [m/s]
    pub a: Vector3<f64>,   // current Euclidean XYZ accelleration [m/s^2]
    pub m: f64,            // particle mass [kg]
    pub r: f64,            // radius of the particle [m]
    pub e: f64,            // potential energy [kcal/mol]
    pub q: f64,            // electrical charge [C]
    pub particle_type: SubatomicParticleType,

    data_trace: HashMap<ParticleProperty, RefCell<Vec<Matrix3x1<f64>>>>
}

RefCell acts as an immutable wrapper to a value it holds, which is of type &T. Using its .borrow_mut()or .get_mut() methods you can mutate the value it wraps. It is a great way of having multiple mutable references at the same time (if you are operating in a single-threaded environment).

So now instead of cloning the whole vector, we will simply get a mutable reference to it and append the element in place, which will truly get rid of the expensive clone operation and is still guaranteed to be safe. My first attempt looked like this:

     pub fn log_debug_data(&mut self) {
        match self.data_trace.get_mut(&ParticleProperty::Position) {
            Some(trace) => { 
                let t = trace.get_mut();
                t.push(self.pos);
            },
            None => { self.data_trace.insert(ParticleProperty::Position, RefCell::new(vec![self.pos])); }
        }
    }

If you look closely I ended up getting not just a mutable reference to the underlying vector I want to mutate, but to the parent RefCell as well. Although this works, there’s no need to do that. Switching from HashMap‘s mutable data_trace.get_mut() to immutable data_trace.get() forced me to change the way of getting the internal vector behind the RefCell from trace.get_mut() to trace.borrow_mut(). Based on the documentation for get_mut() this seems to be a very good thing. The final code looks like this:

    pub fn log_debug_data(&mut self) {
        match self.data_trace.get(&ParticleProperty::Position) {
            Some(trace) => { 
                let mut t = trace.borrow_mut();
                t.push(self.pos);
            },
            None => { self.data_trace.insert(ParticleProperty::Position, RefCell::new(vec![self.pos])); }
        }
    }

Leave a Comment