Specializing over literals/enums at compile time

Imagine a function taking a: f64 and b: f64 as an input.

I want to express the fact that if b == 0.0 (at compile time), a+b == a, therefore building two variants of the function, one for b == 0.0 (with the optimization) and one for any other b (with the full sum).

One obvious choice is defining a wrapper enum type (for the sake of the argument, I could simply reuse Option<f64>) where some_b == None if b == 0.0 and some_b = Some(b) otherwise and use a match block.

Another choice is to define two different types implementing a new trait, transforming what would be a branch selection in the aforementioned enum solution by dynamic dispatch.

Both those solution are costlier than just copy-pasting my code and making the optimisations by hand in one variant. it also prevents vectorisation from happening (I chose floats and addition for simplicity, the motivating driver is not that).


Would that something that could only be done with macros or is there something obvious in the type system I did not consider?

The usual solution is to clearly separate the function into one or more fast paths and slow paths (e.g. one with a == 0, one with b == 0 and the general case), then rely on downstream optimizations such as inlining and const propagation to emit more efficient code when the arguments are known constants at compile time.

1 Like

Note that if you use generic functions instead of dyn Traits, Rust will use static dispatch. Combined with #[inline(always)], the optimizer can be convinced to produce some very efficient code. For example, code like this will let the compiler optimize out all the statically-known Zeros as constant values:

trait UsizeVal: Copy+'static { fn get(self)->usize; }
impl UsizeVal for usize {
    #[inline(always)]
    fn get(self)->usize { self )
}

#[derive(Copy,Clone)] struct Zero;
impl UsizeVal for Zero {
    #[inline(always)]
    fn get(self)->usize { 0 }
}
1 Like

Do you mean that I should create a small function shim with a few explicit conditionals/matches calling specific implementations? And the compiler should remove them as long as the branch can be resolved at compile time?

Ohhh of course! Made the wrong assumption here about traits... Just assumed everything was dyn!

Yes.

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.