Eyes Above The Waves

Robert O'Callahan. Christian. Repatriate Kiwi. Hacker.

Sunday 19 January 2020

Static Customization Of Function Signatures In Rust

Sometimes I have a big function that does a lot, and in new code I need to do almost the same thing, but slightly differently. Often the best approach is to call the same function but add parameters (often described as "options") to select specific variations. This can get ugly when the function has optional outputs — results that are only produced when certain options were passed in — because typically there is the possibility of an error when code looks for an output (e.g. unwraps a Rust Option) at a site that did not request it. It would be great if the compiler could check that you only use an output when you passed in the option to enable it. Fortunately, some simple Rust coding patterns let us achieve this.

Here's an example simplified from some Pernosco code. We have a function that pretty-prints a data structure at a specific moment in time, which takes an optional parameter specifying another time to render the change in values between the two times. If that optional parameter is specified, the function returns an additional result — whether or not there was any difference in the two values. The naive signature would look something like this:

fn print(&self, moment: Moment, other: Option<Moment>)
-> (String, Option<Difference>) {
  ...
  let difference = other.map(|moment| ...);
  ...
  (..., difference)
}
...
let (string1, difference1) = value.print(moment, None);
let (string2, difference2) = value.print(moment, Some(other_moment));
println!("{:?}", difference1.unwrap()); // PANIC
println!("{:?}", difference2.unwrap());
It is possible to misuse this function by passing None in other and then expecting to find a meaningful value in the second part of the result pair. We'd probably catch it in tests, but it would be good to catch it at compile time. There is also a small efficiency issue: passing None for other is a bit less efficient than calling a customized version of print that has been optimized to remove other and the Difference result.

Here's a way to avoid those problems:

trait PrintOptionalMoment {
  type PrintOptionalDifference;
  fn map<F>(&self, closure: F) -> Self::PrintOptionalDifference
     where F: FnOnce(Moment) -> Difference;
}
impl PrintOptionalMoment for () {
  type PrintOptionalDifference = ();
  fn map<F>(&self, closure: F) -> Self::PrintOptionalDifference
     where F: FnOnce(Moment) -> Difference {}
}
impl PrintOptionalMoment for Moment {
  type PrintOptionalDifference = Difference;
  fn map<F>(&self, closure: F) -> Self::PrintOptionalDifference
     where F: FnOnce(Moment) -> Difference { closure(*self) }
}
fn print<Opt: PrintOptionalMoment>(&self, moment: Moment, other: Opt)
-> (String, Opt::PrintOptionalDifference) {
  ...
  let difference = other.map(|moment| ...);
  ...
  (..., difference)
}
let (string1, ()) = value.print(moment, ());
let (string2, difference2) = value.print(moment, other_moment);
println!("{:?}", difference2);

Rust playground link.

This cleans up the call sites nicely. When you don't pass other, the "difference" result has type (), so you can't misuse it or cause a panic trying to unwrap it. When you do pass other, the "difference" result is not an Option, so you don't need to unwrap it. The implementation of print is basically unchanged, but now Rust will generate two versions of the function, and the version that doesn't take other should be optimized about as well as a handwritten function that removed other. (Unlike in C++, in Rust a () value does not take any space in a struct or tuple.)

If you're not familiar with Rust, the intuition here is that we define a trait PrintOptionalMoment that we use to mean "a type that is either nothing, or a Moment", and we declare that the "void" type () and the type Moment both satisfy PrintOptionalMoment. Then we make print generic over type Opt, which can be either of those. The PrintOptionalMoment trait defines an associated type PrintOptionalDifference which is the result type associated with each Opt that satisfies PrintOptionalMoment.

This approach easily extends to cover more complicated relationships between options and input and output types. In some situations it might be more trouble than it's worth, or the generated code duplication is undesirable, but I think it's a good tool to have.