A simple fast path for truncate_str code when the entire string is ASCII. Related PR: https://github.com/ClementTsang/bottom/pull/1330

truncate_str_bench.rs

//! ```cargo
//! [dependencies]
//! unicode-segmentation = "1.10.1"
//! unicode-width = "0.1.11"
//!
//! [dev-dependencies]
//! criterion = "0.5"
//!
//! [[bench]]
//! name = "truncate_str_bench"
//! harness = false
//! ```

use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;

#[inline]
fn grapheme_width(g: &str) -> usize {
    if g.contains('\u{200d}') {
        2
    } else {
        UnicodeWidthStr::width(g)
    }
}

#[inline]
fn fast_ascii<U: Into<usize>>(content: &str, width: U) -> String {
    let width = width.into();

    if width > 0 {
        if content.is_ascii() {
            // If the entire string is ASCII, we can use a much simpler approach.

            if content.len() <= width {
                content.to_owned()
            } else {
                let mut text = String::with_capacity(width);
                let (keep, _throw) = content.split_at(width - 1);
                text.push_str(keep);
                text.push('…');

                text
            }
        } else {
            let mut text = String::with_capacity(width);
            let mut curr_width = 0;
            let mut early_break = false;

            // This tracks the length of the last added string - note this does NOT match the grapheme *width*.
            let mut last_added_str_len = 0;

            // Cases to handle:
            // - Completes adding the entire string.
            // - Adds a character up to the boundary, then fails.
            // - Adds a character not up to the boundary, then fails.
            // Inspired by https://tomdebruijn.com/posts/rust-string-length-width-calculations/
            for g in UnicodeSegmentation::graphemes(content, true) {
                let g_width = grapheme_width(g);

                if curr_width + g_width <= width {
                    curr_width += g_width;
                    last_added_str_len = g.len();
                    text.push_str(g);
                } else {
                    early_break = true;
                    break;
                }
            }

            if early_break {
                if curr_width == width {
                    // Remove the last grapheme cluster added.
                    text.truncate(text.len() - last_added_str_len);
                }
                text.push('…');
            }
            text
        }
    } else {
        String::default()
    }
}

#[inline]
fn original<U: Into<usize>>(content: &str, width: U) -> String {
    let width = width.into();
    let mut text = String::with_capacity(width);

    if width > 0 {
        let mut curr_width = 0;
        let mut early_break = false;

        // This tracks the length of the last added string - note this does NOT match the grapheme *width*.
        let mut last_added_str_len = 0;

        // Cases to handle:
        // - Completes adding the entire string.
        // - Adds a character up to the boundary, then fails.
        // - Adds a character not up to the boundary, then fails.
        // Inspired by https://tomdebruijn.com/posts/rust-string-length-width-calculations/
        for g in UnicodeSegmentation::graphemes(content, true) {
            let g_width = grapheme_width(g);

            if curr_width + g_width <= width {
                curr_width += g_width;
                last_added_str_len = g.len();
                text.push_str(g);
            } else {
                early_break = true;
                break;
            }
        }

        if early_break {
            if curr_width == width {
                // Remove the last grapheme cluster added.
                text.truncate(text.len() - last_added_str_len);
            }
            text.push('…');
        }
    }

    text
}

fn ascii_only(c: &mut Criterion) {
    let mut group = c.benchmark_group("ascii only");
    let content = "012345678910";

    for i in [1_usize, 5, 10, 15].iter() {
        group.bench_with_input(BenchmarkId::new("original", i), i, |b, i| {
            b.iter(|| original(black_box(content), *i))
        });

        group.bench_with_input(BenchmarkId::new("fast_ascii", i), i, |b, i| {
            b.iter(|| fast_ascii(black_box(content), *i))
        });
    }

    group.finish();
}

fn header(c: &mut Criterion) {
    let mut group = c.benchmark_group("bottom header example");
    let content = "CPU(c)▲";

    for i in [1_usize, 2, 3, 4, 5, 6, 7, 8, 9, 10].iter() {
        group.bench_with_input(BenchmarkId::new("original", i), i, |b, i| {
            b.iter(|| original(black_box(content), *i))
        });

        group.bench_with_input(BenchmarkId::new("fast_ascii", i), i, |b, i| {
            b.iter(|| fast_ascii(black_box(content), *i))
        });
    }

    group.finish();
}

fn unicode_only(c: &mut Criterion) {
    let mut group = c.benchmark_group("unicode only");
    let content = "施氏食獅施氏食獅";

    for i in [1_usize, 5, 10, 20].iter() {
        group.bench_with_input(BenchmarkId::new("original", i), i, |b, i| {
            b.iter(|| original(black_box(content), *i))
        });

        group.bench_with_input(BenchmarkId::new("fast_ascii", i), i, |b, i| {
            b.iter(|| fast_ascii(black_box(content), *i))
        });
    }

    group.finish();
}

fn ascii_unicode_mix(c: &mut Criterion) {
    let mut group = c.benchmark_group("ascii unicode mix");
    let content = "abc施氏食獅abc";

    for i in [1_usize, 5, 10, 20].iter() {
        group.bench_with_input(BenchmarkId::new("original", i), i, |b, i| {
            b.iter(|| original(black_box(content), *i))
        });

        group.bench_with_input(BenchmarkId::new("fast_ascii", i), i, |b, i| {
            b.iter(|| fast_ascii(black_box(content), *i))
        });
    }

    group.finish();
}

criterion_group!(benches, ascii_only, header, unicode_only, ascii_unicode_mix);
criterion_main!(benches);

/*
Results running `cargo criterion` (Ryzen 7 5700X, 32GiB RAM, Arch Linux 6.6.1):

ascii only/original/1   time:   [56.157 ns 56.539 ns 57.029 ns]
ascii only/fast_ascii/1 time:   [23.273 ns 23.364 ns 23.487 ns]
ascii only/original/5   time:   [113.56 ns 113.90 ns 114.27 ns]
ascii only/fast_ascii/5 time:   [20.109 ns 20.175 ns 20.228 ns]
ascii only/original/10  time:   [179.32 ns 179.66 ns 180.21 ns]
ascii only/fast_ascii/10
                        time:   [19.607 ns 19.620 ns 19.633 ns]
ascii only/original/15  time:   [188.93 ns 189.24 ns 189.50 ns]
ascii only/fast_ascii/15
                        time:   [13.407 ns 13.422 ns 13.450 ns]

bottom header example/original/1
                        time:   [55.596 ns 55.666 ns 55.748 ns]
bottom header example/fast_ascii/1
                        time:   [56.849 ns 56.890 ns 56.934 ns]
bottom header example/original/2
                        time:   [69.841 ns 70.205 ns 70.862 ns]
bottom header example/fast_ascii/2
                        time:   [70.885 ns 71.126 ns 71.559 ns]
bottom header example/original/3
                        time:   [82.544 ns 82.652 ns 82.751 ns]
bottom header example/fast_ascii/3
                        time:   [83.702 ns 83.815 ns 83.930 ns]
bottom header example/original/4
                        time:   [97.009 ns 97.120 ns 97.237 ns]
bottom header example/fast_ascii/4
                        time:   [96.921 ns 97.015 ns 97.114 ns]
bottom header example/original/5
                        time:   [111.65 ns 111.91 ns 112.28 ns]
bottom header example/fast_ascii/5
                        time:   [113.44 ns 114.05 ns 115.17 ns]
bottom header example/original/6
                        time:   [129.86 ns 130.11 ns 130.47 ns]
bottom header example/fast_ascii/6
                        time:   [130.50 ns 131.22 ns 131.79 ns]
bottom header example/original/7
                        time:   [131.41 ns 131.84 ns 132.20 ns]
bottom header example/fast_ascii/7
                        time:   [135.24 ns 135.35 ns 135.47 ns]
bottom header example/original/8
                        time:   [128.13 ns 128.43 ns 128.83 ns]
bottom header example/fast_ascii/8
                        time:   [134.53 ns 134.68 ns 134.85 ns]
bottom header example/original/9
                        time:   [126.76 ns 126.94 ns 127.15 ns]
bottom header example/fast_ascii/9
                        time:   [126.28 ns 126.57 ns 127.00 ns]
bottom header example/original/10
                        time:   [127.37 ns 127.91 ns 128.37 ns]
bottom header example/fast_ascii/10
                        time:   [122.80 ns 124.35 ns 126.61 ns]

unicode only/original/1 time:   [56.105 ns 56.200 ns 56.302 ns]
unicode only/fast_ascii/1
                        time:   [56.408 ns 56.463 ns 56.519 ns]
unicode only/original/5 time:   [84.313 ns 84.413 ns 84.520 ns]
unicode only/fast_ascii/5
                        time:   [83.555 ns 83.731 ns 83.957 ns]
unicode only/original/10
                        time:   [131.49 ns 131.59 ns 131.69 ns]
unicode only/fast_ascii/10
                        time:   [130.64 ns 131.28 ns 132.12 ns]
unicode only/original/20
                        time:   [184.25 ns 184.59 ns 185.07 ns]
unicode only/fast_ascii/20
                        time:   [182.64 ns 182.73 ns 182.83 ns]

ascii unicode mix/original/1
                        time:   [55.785 ns 55.889 ns 56.001 ns]
ascii unicode mix/fast_ascii/1
                        time:   [56.824 ns 56.908 ns 56.992 ns]
ascii unicode mix/original/5
                        time:   [104.29 ns 104.41 ns 104.55 ns]
ascii unicode mix/fast_ascii/5
                        time:   [101.65 ns 101.73 ns 101.82 ns]
ascii unicode mix/original/10
                        time:   [135.70 ns 135.79 ns 135.91 ns]
ascii unicode mix/fast_ascii/10
                        time:   [133.93 ns 135.05 ns 137.08 ns]
ascii unicode mix/original/20
                        time:   [173.62 ns 174.01 ns 174.48 ns]
ascii unicode mix/fast_ascii/20
                        time:   [171.50 ns 172.04 ns 172.81 ns]
*/