Skip to content

Commit

Permalink
Auto merge of rust-lang#125679 - clarfonthey:escape_ascii, r=joboet
Browse files Browse the repository at this point in the history 
Optimize `escape_ascii` using a lookup table

Based upon my suggestion here: rust-lang#125340 (comment)

Effectively, we can take advantage of the fact that ASCII only needs 7 bits to make the eighth bit store whether the value should be escaped or not. This adds a 256-byte lookup table, but 256 bytes *should* be small enough that very few people will mind, according to my probably not incontrovertible opinion.

The generated assembly isn't clearly better (although has fewer branches), so, I decided to benchmark on three inputs: first on a random 200KiB, then on `/bin/cat`, then on `Cargo.toml` for this repo. In all cases, the generated code ran faster on my machine. (an old i7-8700)

But, if you want to try my benchmarking code for yourself:

<details><summary>Criterion code below. Replace <code>/home/ltdk/rustsrc</code> with the appropriate directory.</summary>

```rust
#![feature(ascii_char)]
#![feature(ascii_char_variants)]
#![feature(const_option)]
#![feature(let_chains)]
use core::ascii;
use core::ops::Range;
use criterion::{criterion_group, criterion_main, Criterion};
use rand::{thread_rng, Rng};

const HEX_DIGITS: [ascii::Char; 16] = *b"0123456789abcdef".as_ascii().unwrap();

#[inline]
const fn backslash<const N: usize>(a: ascii::Char) -> ([ascii::Char; N], Range<u8>) {
    const { assert!(N >= 2) };

    let mut output = [ascii::Char::Null; N];

    output[0] = ascii::Char::ReverseSolidus;
    output[1] = a;

    (output, 0..2)
}

#[inline]
const fn hex_escape<const N: usize>(byte: u8) -> ([ascii::Char; N], Range<u8>) {
    const { assert!(N >= 4) };

    let mut output = [ascii::Char::Null; N];

    let hi = HEX_DIGITS[(byte >> 4) as usize];
    let lo = HEX_DIGITS[(byte & 0xf) as usize];

    output[0] = ascii::Char::ReverseSolidus;
    output[1] = ascii::Char::SmallX;
    output[2] = hi;
    output[3] = lo;

    (output, 0..4)
}

#[inline]
const fn verbatim<const N: usize>(a: ascii::Char) -> ([ascii::Char; N], Range<u8>) {
    const { assert!(N >= 1) };

    let mut output = [ascii::Char::Null; N];

    output[0] = a;

    (output, 0..1)
}

/// Escapes an ASCII character.
///
/// Returns a buffer and the length of the escaped representation.
const fn escape_ascii_old<const N: usize>(byte: u8) -> ([ascii::Char; N], Range<u8>) {
    const { assert!(N >= 4) };

    match byte {
        b'\t' => backslash(ascii::Char::SmallT),
        b'\r' => backslash(ascii::Char::SmallR),
        b'\n' => backslash(ascii::Char::SmallN),
        b'\\' => backslash(ascii::Char::ReverseSolidus),
        b'\'' => backslash(ascii::Char::Apostrophe),
        b'\"' => backslash(ascii::Char::QuotationMark),
        0x00..=0x1F => hex_escape(byte),
        _ => match ascii::Char::from_u8(byte) {
            Some(a) => verbatim(a),
            None => hex_escape(byte),
        },
    }
}

/// Escapes an ASCII character.
///
/// Returns a buffer and the length of the escaped representation.
const fn escape_ascii_new<const N: usize>(byte: u8) -> ([ascii::Char; N], Range<u8>) {
    /// Lookup table helps us determine how to display character.
    ///
    /// Since ASCII characters will always be 7 bits, we can exploit this to store the 8th bit to
    /// indicate whether the result is escaped or unescaped.
    ///
    /// We additionally use 0x80 (escaped NUL character) to indicate hex-escaped bytes, since
    /// escaped NUL will not occur.
    const LOOKUP: [u8; 256] = {
        let mut arr = [0; 256];
        let mut idx = 0;
        loop {
            arr[idx as usize] = match idx {
                // use 8th bit to indicate escaped
                b'\t' => 0x80 | b't',
                b'\r' => 0x80 | b'r',
                b'\n' => 0x80 | b'n',
                b'\\' => 0x80 | b'\\',
                b'\'' => 0x80 | b'\'',
                b'"' => 0x80 | b'"',

                // use NUL to indicate hex-escaped
                0x00..=0x1F | 0x7F..=0xFF => 0x80 | b'\0',

                _ => idx,
            };
            if idx == 255 {
                break;
            }
            idx += 1;
        }
        arr
    };

    let lookup = LOOKUP[byte as usize];

    // 8th bit indicates escape
    let lookup_escaped = lookup & 0x80 != 0;

    // SAFETY: We explicitly mask out the eighth bit to get a 7-bit ASCII character.
    let lookup_ascii = unsafe { ascii::Char::from_u8_unchecked(lookup & 0x7F) };

    if lookup_escaped {
        // NUL indicates hex-escaped
        if matches!(lookup_ascii, ascii::Char::Null) {
            hex_escape(byte)
        } else {
            backslash(lookup_ascii)
        }
    } else {
        verbatim(lookup_ascii)
    }
}

fn escape_bytes(bytes: &[u8], f: impl Fn(u8) -> ([ascii::Char; 4], Range<u8>)) -> Vec<ascii::Char> {
    let mut vec = Vec::new();
    for b in bytes {
        let (buf, range) = f(*b);
        vec.extend_from_slice(&buf[range.start as usize..range.end as usize]);
    }
    vec
}

pub fn criterion_benchmark(c: &mut Criterion) {
    let mut group = c.benchmark_group("escape_ascii");

    group.sample_size(1000);

    let rand_200k = &mut [0; 200 * 1024];
    thread_rng().fill(&mut rand_200k[..]);
    let cat = include_bytes!("/bin/cat");
    let cargo_toml = include_bytes!("/home/ltdk/rustsrc/Cargo.toml");

    group.bench_function("old_rand", |b| {
        b.iter(|| escape_bytes(rand_200k, escape_ascii_old));
    });
    group.bench_function("new_rand", |b| {
        b.iter(|| escape_bytes(rand_200k, escape_ascii_new));
    });

    group.bench_function("old_bin", |b| {
        b.iter(|| escape_bytes(cat, escape_ascii_old));
    });
    group.bench_function("new_bin", |b| {
        b.iter(|| escape_bytes(cat, escape_ascii_new));
    });

    group.bench_function("old_cargo_toml", |b| {
        b.iter(|| escape_bytes(cargo_toml, escape_ascii_old));
    });
    group.bench_function("new_cargo_toml", |b| {
        b.iter(|| escape_bytes(cargo_toml, escape_ascii_new));
    });

    group.finish();
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);
```

</details>

My benchmark results:

```
escape_ascii/old_rand   time:   [1.6965 ms 1.7006 ms 1.7053 ms]
Found 22 outliers among 1000 measurements (2.20%)
  4 (0.40%) high mild
  18 (1.80%) high severe
escape_ascii/new_rand   time:   [1.6749 ms 1.6953 ms 1.7158 ms]
Found 38 outliers among 1000 measurements (3.80%)
  38 (3.80%) high mild
escape_ascii/old_bin    time:   [224.59 µs 225.40 µs 226.33 µs]
Found 39 outliers among 1000 measurements (3.90%)
  17 (1.70%) high mild
  22 (2.20%) high severe
escape_ascii/new_bin    time:   [164.86 µs 165.63 µs 166.58 µs]
Found 107 outliers among 1000 measurements (10.70%)
  43 (4.30%) high mild
  64 (6.40%) high severe
escape_ascii/old_cargo_toml
                        time:   [23.397 µs 23.699 µs 24.014 µs]
Found 204 outliers among 1000 measurements (20.40%)
  21 (2.10%) high mild
  183 (18.30%) high severe
escape_ascii/new_cargo_toml
                        time:   [16.404 µs 16.438 µs 16.483 µs]
Found 88 outliers among 1000 measurements (8.80%)
  56 (5.60%) high mild
  32 (3.20%) high severe
```

Random: 1.7006ms => 1.6953ms (<1% speedup)
Binary: 225.40µs => 165.63µs (26% speedup)
Text: 23.699µs => 16.438µs (30% speedup)
  • Loading branch information
@bors
bors committed Oct 13, 2024
2 parents 1799481 + 86c7526 commit 9c1e162
Show file tree
Hide file tree
Showing 2 changed files with 110 additions and 26 deletions.
114 changes: 91 additions & 23 deletions core/src/escape.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,38 +18,106 @@ const fn backslash<const N: usize>(a: ascii::Char) -> ([ascii::Char; N], Range<u
(output, 0..2)
}

#[inline]
const fn hex_escape<const N: usize>(byte: u8) -> ([ascii::Char; N], Range<u8>) {
const { assert!(N >= 4) };

let mut output = [ascii::Char::Null; N];

let hi = HEX_DIGITS[(byte >> 4) as usize];
let lo = HEX_DIGITS[(byte & 0xf) as usize];

output[0] = ascii::Char::ReverseSolidus;
output[1] = ascii::Char::SmallX;
output[2] = hi;
output[3] = lo;

(output, 0..4)
}

#[inline]
const fn verbatim<const N: usize>(a: ascii::Char) -> ([ascii::Char; N], Range<u8>) {
const { assert!(N >= 1) };

let mut output = [ascii::Char::Null; N];

output[0] = a;

(output, 0..1)
}

/// Escapes an ASCII character.
///
/// Returns a buffer and the length of the escaped representation.
const fn escape_ascii<const N: usize>(byte: u8) -> ([ascii::Char; N], Range<u8>) {
const { assert!(N >= 4) };

match byte {
b'\t' => backslash(ascii::Char::SmallT),
b'\r' => backslash(ascii::Char::SmallR),
b'\n' => backslash(ascii::Char::SmallN),
b'\\' => backslash(ascii::Char::ReverseSolidus),
b'\'' => backslash(ascii::Char::Apostrophe),
b'\"' => backslash(ascii::Char::QuotationMark),
byte => {
let mut output = [ascii::Char::Null; N];

if let Some(c) = byte.as_ascii()
&& !byte.is_ascii_control()
{
output[0] = c;
(output, 0..1)
} else {
let hi = HEX_DIGITS[(byte >> 4) as usize];
let lo = HEX_DIGITS[(byte & 0xf) as usize];
#[cfg(feature = "optimize_for_size")]
{
match byte {
b'\t' => backslash(ascii::Char::SmallT),
b'\r' => backslash(ascii::Char::SmallR),
b'\n' => backslash(ascii::Char::SmallN),
b'\\' => backslash(ascii::Char::ReverseSolidus),
b'\'' => backslash(ascii::Char::Apostrophe),
b'"' => backslash(ascii::Char::QuotationMark),
0x00..=0x1F | 0x7F => hex_escape(byte),
_ => match ascii::Char::from_u8(byte) {
Some(a) => verbatim(a),
None => hex_escape(byte),
},
}
}

#[cfg(not(feature = "optimize_for_size"))]
{
/// Lookup table helps us determine how to display character.
///
/// Since ASCII characters will always be 7 bits, we can exploit this to store the 8th bit to
/// indicate whether the result is escaped or unescaped.
///
/// We additionally use 0x80 (escaped NUL character) to indicate hex-escaped bytes, since
/// escaped NUL will not occur.
const LOOKUP: [u8; 256] = {
let mut arr = [0; 256];
let mut idx = 0;
while idx <= 255 {
arr[idx] = match idx as u8 {
// use 8th bit to indicate escaped
b'\t' => 0x80 | b't',
b'\r' => 0x80 | b'r',
b'\n' => 0x80 | b'n',
b'\\' => 0x80 | b'\\',
b'\'' => 0x80 | b'\'',
b'"' => 0x80 | b'"',

// use NUL to indicate hex-escaped
0x00..=0x1F | 0x7F..=0xFF => 0x80 | b'\0',

idx => idx,
};
idx += 1;
}
arr
};

output[0] = ascii::Char::ReverseSolidus;
output[1] = ascii::Char::SmallX;
output[2] = hi;
output[3] = lo;
let lookup = LOOKUP[byte as usize];

(output, 0..4)
// 8th bit indicates escape
let lookup_escaped = lookup & 0x80 != 0;

// SAFETY: We explicitly mask out the eighth bit to get a 7-bit ASCII character.
let lookup_ascii = unsafe { ascii::Char::from_u8_unchecked(lookup & 0x7F) };

if lookup_escaped {
// NUL indicates hex-escaped
if matches!(lookup_ascii, ascii::Char::Null) {
hex_escape(byte)
} else {
backslash(lookup_ascii)
}
} else {
verbatim(lookup_ascii)
}
}
}
Expand Down
22 changes: 19 additions & 3 deletions core/tests/ascii.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -481,9 +481,25 @@ fn ascii_ctype_const() {
}

#[test]
fn test_ascii_display() {
assert_eq!(b"foo'bar".escape_ascii().to_string(), r#"foo\'bar"#);
assert_eq!(b"\0\xff".escape_ascii().to_string(), r#"\x00\xff"#);
fn test_escape_ascii() {
let mut buf = [0u8; 0x1F + 7]; // 0..=0x1F plus two quotes, slash, \x7F, \x80, \xFF
for idx in 0..=0x1F {
buf[idx] = idx as u8;
}
buf[0x20] = b'\'';
buf[0x21] = b'"';
buf[0x22] = b'\\';
buf[0x23] = 0x7F;
buf[0x24] = 0x80;
buf[0x25] = 0xff;
assert_eq!(
buf.escape_ascii().to_string(),
r#"\x00\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\x0c\r\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1a\x1b\x1c\x1d\x1e\x1f\'\"\\\x7f\x80\xff"#
);
}

#[test]
fn test_escape_ascii_iter() {
let mut it = b"\0fastpath\xffremainder\xff".escape_ascii();
let _ = it.advance_by(4);
let _ = it.advance_back_by(4);
Expand Down

0 comments on commit 9c1e162

Please sign in to comment.