Rust Operator Overloading

2020-03-23 // 5 minutes

In today's blog post, we're going to explore the options for implementing operator overloading in Rust. To start with, let's say we've defined a Vec3f struct, as below. Note the std::ops reference, we'll need that for the actual overloading call later.

use std::ops;

#[allow(dead_code)]
pub struct Vec3f {
	pub x: f32,
	pub y: f32,
	pub z: f32
}

Let's start out with a simple vector operation - addition. In case your math skills are a little rusty, we want two vectors A and B to be added together into vector C such that: C.x = A.x + B.x, C.y = A.y + B.y, and C.z = A.z + B.z. Let's start out with a straightforward solution as illustrated below.

impl ops::Add for Vec3f {
	type Output = Vec3f;

	fn add(self, other: Vec3f) -> Vec3f {
		Vec3f {
			x : self.x + other.x,
			y : self.y + other.y,
			z : self.z + other.z
		}
	}
}

This should be fairly straightforward to follow. The only unusual thing to note is setting type Output equal to Vec3f. The first operand, A, is referred to as self and the second, B, is referred to as other. Let's set up a quick example to verify that this works.

fn main() {
    let a : Vec3f = Vec3f{
    	x : 1.0,
    	y : 2.0,
    	z : 3.0
    };

    let b : Vec3f = Vec3f{
    	x : 2.0,
    	y : 3.0,
    	z : 4.0
    };

    let c : Vec3f = a + b;

    println!("[{}, {}, {}]", c.x, c.y, c.z);
}

If you've followed up until now, your cargo run should spit out [3, 5, 7] for c. So, where does this get complicated? If you were thinking the borrow checker, then you'd be correct. If you need a quick refresher, make sure to check out the official docs on references and borrowing. But as a quick review, if you reference a variable after it's definition, it is considered to be moved. You can't reference it after the move or the compiler will complain - this helps Rust keep track of variable ownership. But what you can do is borrow the variable with the & symbol.

let c : Vec3f = a + b;
let d : Vec3f = c + &a;
println!("[{}, {}, {}]", c.x, c.y, c.z);

For our purposes, let's modify the main function to demonstrate why the borrow checker won't let us get away with just the simple version. As above, change your main to include the definition for the variable d. Notice the & before the a? Once again, we're borrowing the variable here. So what's the problem? Your operator overload expects a Vec3f and you've just provided a &Vec3f.

Rust treats a borrowed variable, &Vec3f in this case, to be a separate type. This makes operator overloading tedious because we now need to define an operator overload for each operation four different times - depending on if the operands are currently borrowed or not. Fret not, there is a better way, but first let's illustrate how tedious this could become for our example.

impl ops::Add for Vec3f {
	type Output = Vec3f;

	fn add(self, other: Vec3f) -> Vec3f {
		Vec3f {
			x : self.x + other.x,
			y : self.y + other.y,
			z : self.z + other.z
		}
	}
}

impl ops::Add for &Vec3f {
	type Output = Vec3f;

	fn add(self, other: Vec3f) -> Vec3f {
		Vec3f {
			x : self.x + other.x,
			y : self.y + other.y,
			z : self.z + other.z
		}
	}
}

impl ops::Add<&Vec3f> for Vec3f {
	type Output = Vec3f;

	fn add(self, other: &Vec3f) -> Vec3f {
		Vec3f {
			x : self.x + other.x,
			y : self.y + other.y,
			z : self.z + other.z
		}
	}
}

impl ops::Add<&Vec3f> for &Vec3f {
	type Output = Vec3f;

	fn add(self, other: &Vec3f) -> Vec3f {
		Vec3f {
			x : self.x + other.x,
			y : self.y + other.y,
			z : self.z + other.z
		}
	}
}

Gross, right? Imagine how quickly that grows once you implement that for every operation for that struct. Thankfully, Rust macros and metaprogramming are here to save the day. However, as it turns out, defining that ourselves involves somewhere between 4 and 5 levels of macro indirection. So, for the purposes of this blog post, we'll use the crate auto_ops that provides us a handy macro shortcut. We'll into the inner workings of this crate in a separate blog post.

use auto_ops::*;
	
impl_op_ex!(+ |a: &Vec3f, b: &Vec3f| -> Vec3f { 
	Vec3f {
		x: a.x + b.x,
		y: a.y + b.y,
		z: a.z + b.z
	}
});

fn main() {
    let a : Vec3f = Vec3f{
    	x : 1.0,
    	y : 2.0,
    	z : 3.0
    };

    let b : Vec3f = Vec3f{
    	x : 2.0,
    	y : 3.0,
    	z : 4.0
    };

    let c : Vec3f = a + &b;
    let d : Vec3f = c + &b;

    println!("[{}, {}, {}]", d.x, d.y, d.z);
}

If you've copied everything correctly, you're println! should display [5, 8, 11]. The impl_op_ex! macro can seem a bit automagical, especially since it's not clear how passing in two borrowed struct types also covers the unborrowed case. However, it it thankfully easy to read and we'll dive into its inner workings in the blog post mentioned above.