Optimise and bug-fix `to_digit`

[daniel-pfeiffer]

Proposal

Problem statement

Implementation of to_digit is convoluted and inefficiently does too much. Yet it still accepts radices 0 & 1. While there is no such system as nullary, unary is a totally different scheme, with only digit 1. That is not implemented by this function. All numbers Rust deals with in any base are in positional notation. There 0 is always the smallest digit. And incrementing the biggest digit gives 10. The smallest radix for which this is possible is two.

Motivating examples or use cases

It is inefficient to reassert the radix for each digit again and again. I propose to split this function into a wrapper to_digit that does its due diligence and a worker to_digit_unchecked. I have checked all callers of to_digit, to see where it is already safe, or can be made safe, to switch to to_digit_unchecked. Maybe each time a comment should be added to the place that guarantees a valid radix, to avert accidentally eliminating the guarantee in future:

bounds already checked outside of loop, so can switch:

• library/core/src/num/mod.rs 2x

literal base in a variable, only valid values, so can switch:

• compiler/rustc_session/src/errors.rs 1x
• rustc_parse/src/lexer/mod.rs 1x

in a loop, should add bounds check and switch:

• library/core/src/net/parser.rs 2x

literal radices, where the compiler can hopefully eliminate the asserts, so no need, but might do it to save compile time:

• compiler/rustc_lexer/src/unescape.rs 4x
• librustdoc/html/render/print_item.rs 6x
• rustc_parse_format/src/lib.rs 1x
• many test files

Solution sketch

I propose three variants to choose the style and assumed efficiency you see fit:

pub const fn to_digit(self, radix: u32) -> Option<u32> {
    assert!(radix >= 2, "to_digit: radix is too low (minimum 2)");
    assert!(radix <= 36, "to_digit: radix is too high (maximum 36)");
    self.to_digit_unchecked(radix)
}

/// ### Correctness
///
/// Callers of this function are responsible that this precondition is satisfied:
///
/// - The radix must be between 2 and 36 inclusive.
///
/// Failing that, the returned value cannot be relied on.
#[inline]
// semi-branchless variant
pub const fn to_digit_unchecked(self, radix: u32) -> Option<u32> {
    let is_digit = (self <= '9') as u32;
    let digit =
        is_digit * // branchless if
            // If not a digit, a number greater than radix will be created.
            (self as u32).wrapping_sub('0' as u32)
        + (1 - is_digit) * // branchless else
            // Force the 6th bit to be set to ensure ascii is lower case.
            (self as u32 | 0b10_0000).wrapping_sub('a' as u32).saturating_add(10);
    // FIXME: once then_some is const fn, use it here
    if digit < radix { Some(digit) } else { None }
}

// conventional variant with match
pub const fn to_digit_unchecked(self, radix: u32) -> Option<u32> {
    let digit =
        match self {
            // If not a digit, a number greater than radix will be created.
            ..='9' => (self as u32).wrapping_sub('0' as u32),

            // Force the 6th bit to be set to ensure ascii is lower case.
            _ => (self as u32 | 0b10_0000).wrapping_sub('a' as u32).saturating_add(10)
        };
    // FIXME: once then_some is const fn, use it here
    if digit < radix { Some(digit) } else { None }
}

// conventional variant with if
pub const fn to_digit_unchecked(self, radix: u32) -> Option<u32> {
    let digit =
        if self <= '9' {
            // If not a digit, a number greater than radix will be created.
            (self as u32).wrapping_sub('0' as u32)
        } else {
            // Force the 6th bit to be set to ensure ascii is lower case.
            (self as u32 | 0b10_0000).wrapping_sub('a' as u32).saturating_add(10)
        };
    // FIXME: once then_some is const fn, use it here
    if digit < radix { Some(digit) } else { None }
}

Links and related work

rust-lang/rust#132428

This discussion claims that the new bounds check might break backwards compatibility. I doubt that code would rely on a broken implementation. But if that is a concern, the new assert could be activated with edition 2024. OTOH, with that rationale one could never fix bugs…

[programmerjake]

other than changing the checks for radix < 2, this seems like premature optimization to me, i expect llvm can already optimize to be even better than the branchless code you gave (multiply can be expensive on some cpus), especially optimizing out the radix <= 36 check.

[programmerjake]

ok, actually not premature optimization, the current to_digit implementation is worse than I thought and generates a big pile of branches. an alternative implementation that compiles to mostly branchless code on x86_64 (when inside a loop so the assert is optimized by llvm to be checked once outside the loop): (edit: added check for radix > 10 to optimize out that part for smaller radixes) https://play.rust-lang.org/?version=stable&mode=release&edition=2021&gist=9f0fbb65ed402b2b1ada5c33c8e82580

pub fn to_digit(digit: char, radix: u32) -> Option<u32> {
    assert!(radix >= 2 && radix <= 36);
    let lc_digit = digit as u32 | 0x20;
    let digit = if digit > '9' && radix > 10 {
        lc_digit.wrapping_sub('a' as u32).wrapping_add(10)
    } else {
        (digit as u32).wrapping_sub('0' as u32)
    };
    if digit < radix {
        Some(digit)
    } else {
        None
    }
}

[lolbinarycat]

as far as I know, this would be the only instance of a safe method containing "unchecked" in its name?

what about a private to_digit_no_radix_check that contains all the logic except the assert, then put #[inline(always)] on to_digit, so the assertion always gets folded into the surrounding code, either removed after being proven impossible (the usual case where radix is a constant), or lifted to the top of a loop.

[lolbinarycat]

also, "unchecked" also feels wrong because it still contains the digit < radix check. we might in the future want a variation that also removes that check, and we'll want to reserve the "unchecked" name for that variant.

[daniel-pfeiffer]

The function has a simple check of “digit or letter?” as spaghetti code. I needed to look three times to exactly get it. I am simplifying that. IMHO that’s a win on its own.

I am not experienced with branchless, which is why I offer three alternate simplifications. If none is a clear winner, they might need to be benchmarked on various CPUs. That’s not something I’m in a position to do. OTOH, if the compiler turns the latter two into the same as what we had, just choose the prettiest!

I’m also not an expert on compiler optimisation. If it can really eliminate an inlined check many statements after a similar one (near the start of from_str_radix,) then splitting the function is not needed. But if splitting is benficial, any suggestive name will do.

@programmerjake As for your new implementation, it’s essentially the same as my 3rd variant. Except you added an unnecessary && radix > 10, which is already a side effect of wrapping_sub with the last if. I doubt the compiler can think that deeply. If the existing comments, which I kept, didn’t convey that clearly, they may need to be improved!

[programmerjake]

Except you added an unnecessary && radix > 10

that's necessary to eliminate the check for digit > '9' when radix is a known constant <= 10. llvm isn't smart enough to remove that part of the code otherwise.

[joshtriplett]

We talked about this in today's libs-api meeting.

We would love a PR for the bugfix to check for radix >= 2, and we'd love a PR switching to the branchless version that LLVM can optimize better.

Beyond that, adding annotations and internal functions to make sure LLVM can properly inline this and hoist the checks would be fine as well.

We'd prefer not to make that internal version public unless there are still substantial performance differences after those PRs; we'd want to see specific performance benchmarks motivating any further change.

[programmerjake]

PR switching to the branchless version that LLVM can optimize better.

which version do you mean by that? the first, second, or third one in the top comment, or the one I wrote?

[joshtriplett]

Whatever improves on the current version and solves the problem; not judging which of the implementations is best.

[daniel-pfeiffer]

@programmerjake I’m not a contributor and wouldn’t currently have time to get in that deep. So I’m happy to let you do the PR. However my 3rd version retains the comments, so it would be a better basis. I like your idea of pulling out the lc into a variable. It would make it clearer, if you had not pulled it out of its if-branch. That again unnecessarily gives a reader something to contemplate…