Refinement in Rust: optimization, arithmetic, and stateful predicates

ยท 1209 words ยท 6 minute read

v0.0.4 of refined, my library providing support for refinement types in Rust, has been released


In my last post, I announced the initial release of refined, a library providing support for refinement types in Rust. Since the release, a number of interesting ideas have been contributed by the community. v0.0.4 of the library has now been released, including support for three important new features along with a number of minor improvements.

The major features that this post will focus on are:

  • optimized: a feature flag that allows runtime bounds checking to be elided via use of an unsafe optimization
  • arithmetic: a feature flag that allows relevant specializations of the Refinement struct to be used directly with the arithmetic traits from std::ops
  • Stateful predicates: the ability for users to define StatefulPredicate implementations, allowing for more efficient refinement in cases where applying predicate logic is both pure and computationally intensive

Optimized ๐Ÿ”—

Shortly after the initial release of refined, an issue was raised asking about the possibility of eliding bounds checks through the use of hint::assert_unchecked. Given the design of refined and the guarantees that it provides, this seemed like an easy win for performance for those who are willing to opt in to unsafe behavior.

When you turn on the optimized feature, you should notice no difference other than that the generated code is faster than it was before. To better understand the impact of this change, we can look at the generated assembly from a small example program using cargo-show-asm.

First, the code that we’ll be inspecting:

 1use refined::{boundable::unsigned::LessThan, prelude::*};
 2
 3type Size = Refinement<usize, LessThan<12>>;
 4
 5#[no_mangle]
 6fn month_name(s: &Size) -> &'static str {
 7    const MONTHS: [&str; 12] = [
 8        "January",
 9        "February",
10        "March",
11        "April",
12        "May",
13        "June",
14        "July",
15        "August",
16        "September",
17        "October",
18        "November",
19        "December",
20    ];
21    MONTHS[**s]
22}
23
24fn main() {
25    let raw_index = std::env::args()
26        .next_back()
27        .unwrap()
28        .parse::<usize>()
29        .unwrap();
30    let index = Size::refine(raw_index).unwrap();
31
32    println!("{}", month_name(&index));
33}

In “standard” Rust (using a u8 as the argument to month_name rather than the Size refinement), a bounds check on the array access is required. Looking at the generated assembly without optimized, we can see that refined shares this behavior:

1$ cargo asm -p refined-optimized --bin refined-optimized month_name
2   Finished `release` profile [optimized] target(s) in 0.01s
 1.section .text.month_name,"ax",@progbits
 2        .globl  month_name
 3        .p2align        4
 4.type   month_name,@function
 5month_name:
 6        .cfi_startproc
 7        mov rdi, qword ptr [rdi]
 8        cmp rdi, 11
 9        ja .LBB11_2
10        shl rdi, 4
11        lea rcx, [rip + .L__unnamed_9]
12        mov rax, qword ptr [rdi + rcx]
13        mov rdx, qword ptr [rdi + rcx + 8]
14        ret
15.LBB11_2:
16        push rax
17        .cfi_def_cfa_offset 16
18        lea rdx, [rip + .L__unnamed_10]
19        mov esi, 12
20        call qword ptr [rip + core::panicking::panic_bounds_check@GOTPCREL]

For those unfamiliar with assembly, the bounds check is right around the start of the function; if the input argument is greater than 11, we drop into a panic handler:

 1...
 2        mov rdi, qword ptr [rdi]
 3        cmp rdi, 11
 4        ja .LBB11_2
 5...
 6.LBB11_2:
 7        push rax
 8        .cfi_def_cfa_offset 16
 9        lea rdx, [rip + .L__unnamed_10]
10        mov esi, 12
11        call qword ptr [rip + core::panicking::panic_bounds_check@GOTPCREL]

Taking a look at enabling optimized, we can see that this bounds check is no longer generated:

1$ cargo asm -p refined-optimized --bin refined-optimized month_name --features optimized
2    Finished `release` profile [optimized] target(s) in 0.01s
 1.section .text.month_name,"ax",@progbits
 2        .globl  month_name
 3        .p2align        4
 4.type   month_name,@function
 5month_name:
 6        .cfi_startproc
 7        mov rcx, qword ptr [rdi]
 8        shl rcx, 4
 9        lea rdx, [rip + .L__unnamed_9]
10        mov rax, qword ptr [rcx + rdx]
11        mov rdx, qword ptr [rcx + rdx + 8]
12        ret

Because refined knows that the predicate must always hold, the bounds check is entirely spurious. Critically, this optimization removes a branch point from the generated code, which could significantly improve performance if the code path is used infrequently enough that good branch prediction cannot come into play. In any case, removing checks that aren’t required is a Good Thing for performance; as such, once the feature has had time for some testing in the wild to prove its soundness, it will likely become a default feature of refined.

Arithmetic ๐Ÿ”—

Many of the provided predicates in refined are numeric. A particularly common case is the desire to constrain numeric values to a subset of the possible values inhabited by the type. These bounding constraints create integral ranges, allowing for implementation of interval arithmetic for such types in some situations. This in turn allows refined to provide “refined arithmetic” with proven bounds that are verified at compile time rather than at run time when the arithmetic feature is enabled.

Similar to implication, this feature requires the use of generic_const_exprs, so its stability cannot be guaranteed.

Here’s a straightforward example:

 1#![allow(incomplete_features)]
 2#![feature(generic_const_exprs)]
 3
 4use refined::{prelude::*, boundable::unsigned::ClosedInterval};
 5
 6type Larger = Refinement<u16, ClosedInterval<100, 1000>>;
 7type Smaller = Refinement<u16, ClosedInterval<50, 500>>;
 8
 9fn main() {
10  let larger = Larger::refine(500).unwrap();
11  let smaller = Smaller::refine(100).unwrap();
12  // Type annotation not required; provided for clarity
13  let result: Refinement<u16, ClosedInterval<0, 20>> = larger / smaller;
14  assert_eq!(*result, 5);
15}

I’ve attempted to implement all combinations of interval arithmetic that I believe are sound in all situations, but if you find an arithmetic operation that you think should be supported but is not, please let me know so that we can get it added to the library.

Stateful predicates ๐Ÿ”—

Stateful predicates, as the name suggests, allows for the use of predicates that contain state. This is useful in situations where the “materialization” of a predicate is costly or prone to potential failure due to e.g. a dependency on an external resource.

A motivating example is the Regex predicate. Regular expressions are often costly to compile; if you find yourself in a situation where you need to refine many different values using the same regular expression, your program would previously spend the majority of its time in refinement re-compiling the same regex over and over for each value’s refinement. Stateful refinement solves this problem by allowing us to cache the compiled predicate for re-use:

 1use refined::{prelude::*, string::Regex};
 2
 3type_string!(FiveLowercaseLetters, "^[a-z]{5}$");
 4type FiveLower = Refinement<String, Regex<FiveLowercaseLetters>>;
 5
 6fn main() {
 7    let predicate = Regex::<FiveLowercaseLetters>::default(); // Compiles the regex
 8
 9    // Same regex re-used for each refinement
10    let good1 = FiveLower::refine_with_state(&predicate, "abcde".to_string());
11    assert!(good1.is_ok());
12    let good2 = FiveLower::refine_with_state(&predicate, "zzzzz".to_string());
13    assert!(good2.is_ok());
14    let bad = FiveLower::refine_with_state(&predicate, "AAAAA".to_string());
15    assert!(bad.is_err());
16}

Note that the regular expression functionality is kept behind the regex feature flag to prevent an unnecessary dependency for those who don’t need it, but this does not affect the core stateful functionality (which is always available).

Wrapping up ๐Ÿ”—

v0.0.4 included a few other small changes, including a contribution guide for new contributors to the project. A complete list of all changes can be found in the changelog.

The library is quickly reaching a stable state where I’ll be releasing a major version to indicate ongoing, stable support for the existing functionality as I continue to build it out over time. Please reach out or drop me an issue if you have any ideas for future improvements or additions!

Enjoyed the post?
Please consider following for new post notifications!