Singpolyma

Technical Blog

Rust Factory Without Box (Trait Object)

Posted on

I’ve been playing around a lot with Rust recently and it’s quickly becoming my second-favourite programming language. One of the things I’ve been playing with is some Object Oriented design concepts as they might apply. For example, consider this code:

fn format_year(n: i32) -> String {
	if n == 0 {
		"0 is not a year".to_string()
	} else if n < 0 {
		format!("{} BC", -n)
	} else {
		format!("{} AD", n)
	}
}

While maybe overkill for this small example, let’s go ahead and replace conditional with polymorphism:

fn format_year(n: Box<Year>) -> String {
	format!("{} {}", n.year(), n.era())
}

trait Year {
	fn year(&self) -> u32;
	fn era(&self) -> String;
}

impl Year {
	fn new(n: i32) -> Box<Year> {
		if n == 0 {
			Box::new(YearZero())
		} else if n < 0 {
			Box::new(YearBC(-n as u32))
		} else {
			Box::new(YearAD(n as u32))
		}
	}
}

struct YearZero();

impl Year for YearZero {
	fn year(&self) -> u32 { 0 }
	fn era(&self) -> String { "is not a year".to_string() }
}

struct YearBC(u32);

impl Year for YearBC {
	fn year(&self) -> u32 { self.0 }
	fn era(&self) -> String { "BC".to_string() }
}

struct YearAD(u32);

impl Year for YearAD {
	fn year(&self) -> u32 { self.0 }
	fn era(&self) -> String { "AD".to_string() }
}

This works, and really does seem to mimic the way this kind of design looks in a class-based Object Oriented language. It has a major disadvantage, however: all our objects are on the heap now (which is likely to cause performance issues). In some cases, this can be fixed by using CPS so that the trait objects could be borrowed references instead of boxed, but that’s both ugly and not always an option. One other design might be to use an enum:

fn format_year(n: Year) -> String {
	format!("{} {}", n.year(), n.era())
}

enum Year {
	YearZero,
	YearBC(u32),
	YearAD(u32)
}

impl Year {
	fn new(n: i32) -> Year {
		if n == 0 {
			Year::YearZero
		} else if n < 0 {
			Year::YearBC(-n as u32)
		} else {
			Year::YearAD(n as u32)
		}
	}

	fn year(&self) -> u32 {
		match self {
			YearZero => 0,
			YearBC(y) => y,
			YearAD(y) => y
		}
	}

	fn era(&self) -> u32 {
		match self {
			YearZero => "is not a year".to_string(),
			YearBC(y) => "BC".to_string(),
			YearAD(y) => "AD".to_string()
		}
	}
}

No more heap allocations! While this is obviously analogous, some might claim we haven’t actually “replaced conditional” at all, though we have at least contained the conditionals in a place where a type only knows about itself, and not about other things that might get passed in. Even if you accept adding match arms on self as “extension”, in terms of open/closed this requires a modification to at least the enum and the factory to add a new case, instead of just the factory as with the trait version.

What is it about the enum version that allows us to avoid the boxing? Well, an enum knows what all the possibilities are, and so the compiler can know the size that needs to be reserved to store any one of those. With the trait case, the compiler can’t know how big the infinite world of possibilities that might implement that trait could be, and so cannot know the size to be reserved: we have to defer that to runtime and use a box. However, the factory will always actually return only a known list of trait implementations… can we exploit that to know the size somehow? What if we create an enum of the structs from the trait version and have the factory return that?

enum YearEnum {
	YearZero(YearZero),
	YearBC(YearBC),
	YearAD(YearAD)
}

impl Year {
	fn new(n: i32) -> YearEnum {
		if n == 0 {
			YearEnum::YearZero(YearZero())
		} else if n < 0 {
			YearEnum::YearBC(YearBC(-n as u32))
		} else {
			YearEnum::YearAD(YearAD(n as u32))
		}
	}
}

impl std::ops::Deref for YearEnum {
	type Target = Year;

	fn deref(&self) -> &Self::Target {
		match self {
			YearEnum::YearZero(x) => x,
			YearEnum::YearBC(x) => x,
			YearEnum::YearAD(x) => x
		}
	}
}

The impl std::ops::Deref will allow us to call any method in the Year trait on the enum as returned from the factory, allowing this to effectively act as a trait object, but with no heap allocations! This seems like exactly what we want, but it’s a lot of boilerplate. Luckily, it’s very mechanical so creating a macro to do this for us is fairly easy (and I’ll throw in a bunch of other obvious trait implementations while we’re at it):

macro_rules! trait_enum {
	($trait:ident, $enum:ident, $( $item:ident ) , *) => {
		enum $enum {
			$(
				$item($item),
			)*
		}

		impl std::ops::Deref for $enum {
			type Target = $trait;

			fn deref(&self) -> &Self::Target {
				match self {
					$(
						$enum::$item(x) => x,
					)*
				}
			}
		}

		impl From<$enum> for Box<$trait> {
			fn from(input: $enum) -> Self {
				match input {
					$(
						$enum::$item(x) => Box::new(x),
					)*
				}
			}
		}

		impl<'a> From<&'a $enum> for &'a $trait {
			fn from(input: &'a $enum) -> Self {
				&**input
			}
		}

		impl<'a> AsRef<$trait + 'a> for $enum {
			fn as_ref(&self) -> &($trait + 'a) {
				&**self
			}
		}

		impl<'a> std::borrow::Borrow<$trait + 'a> for $enum {
			fn borrow(&self) -> &($trait + 'a) {
				&**self
			}
		}

		$(
			impl From<$item> for $enum {
				fn from(input: $item) -> Self {
					$enum::$item(input)
				}
			}
		)*
	}
}

And now to repeat the first refactoring, but with the help of this new macro:

fn format_year<Y: Year + ?Sized>(n: &Y) -> String {
	format!("{} {}", n.year(), n.era())
}

trait Year {
	fn year(&self) -> u32;
	fn era(&self) -> String;
}

trait_enum!(Year, YearEnum, YearZero, YearBC, YearAD);

impl Year {
	fn new(n: i32) -> YearEnum {
		if n == 0 {
			YearZero().into()
		} else if n < 0 {
			YearBC(-n as u32).into()
		} else {
			YearAD(n as u32).into()
		}
	}
}

struct YearZero();

impl Year for YearZero {
	fn year(&self) -> u32 { 0 }
	fn era(&self) -> String { "is not a year".to_string() }
}

struct YearBC(u32);

impl Year for YearBC {
	fn year(&self) -> u32 { self.0 }
	fn era(&self) -> String { "BC".to_string() }
}

struct YearAD(u32);

impl Year for YearAD {
	fn year(&self) -> u32 { self.0 }
	fn era(&self) -> String { "AD".to_string() }
}

We do still have two places with must be modified rather than extended (the macro invocation and the factory), but all other code can be written ignorant of those and in the same style as using a normal trait object. The normal trait objects can even be recovered using various implementations the macro creates, or even just by doing &* on the enum. Benchmarking these three styles on a somewhat more complex example actually found this last one to also be the most performant (though only marginally faster than the pure-enum approach), and the boxed-trait-object style to be more than three times slower.

So there you go, next time you ask yourself if you want the flexibility of a trait or the size guarantees and performance of an enum, maybe grab a macro and say: why not both!

2 Responses

Tim Makarios

That’s quite clever!

Just a little typo I spotted near the end of the last code block: you’ve got “impl Year for YearBC” twice instead of “… YearAD” the second time. Just checking further back, it looks like this typo occurs in the first polymorphism example, too.

Thanks for sharing the idea.

Leave a Response