jieba-rs 分词性能优化记录：提升 2.4 倍

2019 年 Paul Meng 写了一篇博客，把 jieba-rs 优化到比 cppjieba 快 33%。之后六年没怎么动过性能，也没人来报慢。

2025 年 5 月开始，有几个贡献者对 hot path 动了手，我也顺手参与了进来。前后一年，七个 PR，HMM 分词从 2.85 µs 降到 1.32 µs，非 HMM 从 2.21 µs 降到 0.94 µs。2.16 倍和 2.35 倍。

单个改动都不算巧妙，大部分看完会觉得"早该这样"。叠在一起效果就很明显了。

怎么测的

每个 commit 单独 checkout，cargo bench，记中位数。机器、输入、allocator 全程不变：

Apple M4 Max
128 GB RAM
macOS 26.4.1
Rust nightly, jemalloc, criterion

测试句子：

我是拖拉机学院手扶拖拉机专业的。不用多久，我就会升职加薪，当上CEO，走上人生巅峰。

典型的中文句子，有标点有英文缩写，搜索引擎索引器喂进来的差不多就这样。

完整数据：

前四个加起来大概 10%。#141 和 #144 各砍掉 25-29%。#145 拿到 9-13%。大头就在后面三个。

lazy_static! → thread_local!，captures → find (PR #122)

Runji Wang

jieba-rs 原来有四个编译好的正则放在 lazy_static! 里，每次访问走一次原子加载来检查初始化。换成 thread_local! 之后变成 TLS 查找，不需要原子操作。

同一个 PR 修了几处本来用 Regex::captures() 但只需要 Regex::find() 的地方。regex crate 文档写得很清楚：is_match > find > captures，每升一级多做一些记录工作。

改动前：

lazy_static! {
    static ref RE_HAN_DEFAULT: Regex = Regex::new(r"...").unwrap();
}
let splitter = SplitMatches::new(&RE_HAN_DEFAULT, sentence);

改动后：

thread_local! {
    static RE_HAN_DEFAULT: Regex = Regex::new(r"...").unwrap();
}
RE_HAN_DEFAULT.with(|re_han| {
    let splitter = SplitMatches::new(re_han, sentence);
    // ...
})

跑出来差 1-2%，和噪声差不多。但做法本身没问题，后面几个 PR 也依赖这个改动。

HMM 工作内存复用 (PR #123)

Runji Wang

Viterbi 解码器每次调用分配三个 Vec：v（概率矩阵），prev（回溯指针），best_path。一个 C 字符的句子配 4 个 HMM 状态，就是三次分配，大小 4*C、4*C、C。单次不大，每个句子都来一遍就累积起来了。

改成 HmmContext 结构体传进来复用：

pub(crate) struct HmmContext {
    v: Vec<f64>,
    prev: Vec<Option<State>>,
    best_path: Vec<State>,
}

clear() + resize() 代替每次重新分配。只有当前句子比之前最长的还长，allocator 才会被调用到。

微基准测试看不出差别，它就一个句子反复跑。实际场景分词几千个不同长度的句子时，省下来的分配开销就有意义了。这也是后面 PR #127 把 context 放到 thread_local 的前提。

借用 slice 代替拷贝 (PR #125)

aniaan

SplitState 枚举有个 into_str() 方法，某些路径上会拷贝匹配到的文本。重构后所有变体直接从输入字符串借用 &str，零拷贝。

改动量不大，lib.rs 里 30 行。但 no_hmm 快了 5-6%，with_hmm 快了 3-4%。内层循环切了上千次字符串，每次少一次拷贝就是这个效果。

thread_local HmmContext 和容量预分配 (PR #127)

PR #123 让 HmmContext 可以复用了，但每次 cut() 还是创建一个新的。这个 PR 放到 thread_local!，同一个线程内跨调用复用：

thread_local! {
    static HMM_CONTEXT: RefCell<HmmContext> = RefCell::new(HmmContext::default());
}

// cut() 里：
HMM_CONTEXT.with(|ctx| {
    let mut hmm_context = ctx.borrow_mut();
    self.cut_dag_hmm(block, &mut words, &mut route, &mut dag, &mut hmm_context);
});

同一个 PR 给 TextRank 和 TfIdf 里几个 Vec::new() 加了 with_capacity()，StaticSparseDAG 的容量乘数从 5 调到 4，加了最小值 32。

cut 基准测试上没看出差别。容量预分配对关键词提取路径更有用，但 benchmark 没覆盖到。

Vec 替代 HashMap，DAG 中打包 word_id (PR #141)

从这个 PR 开始数字动了。所有 cut 模式快了 17-29%，改了四个地方。

最大的一个：StaticSparseDAG 用 HashMap<usize, usize> 存字节偏移到边数组位置的映射。字节偏移就是从 0 开始的连续整数，对这种 key 用 HashMap 完全没必要。直接 Vec<usize> 按下标访问：

// 改动前
start_pos: HashMap<usize, usize>,

// 改动后
start_pos: Vec<usize>,  // 下标 = 字节偏移，值 = 边数组中的位置

空位存 usize::MAX 做 sentinel。clear() 时 fill(NO_ENTRY) 代替 HashMap::clear()。

第二个改动：DAG 的边里直接存 word_id。之前 DAG 只存 (byte_start, byte_end)，到了动态规划 calc() 阶段，每条边再做一次 cedar.exact_match_search() 查词频：

// 改动前：calc() 里重复查 trie
for byte_end in dag.iter_edges(byte_start) {
    let wfrag = &sentence[byte_start..byte_end];
    let freq = if let Some((word_id, _, _)) = self.cedar.exact_match_search(wfrag) {
        self.records[word_id as usize].freq
    } else {
        1
    };
}

改成 DAG 构建时就把 word_id 和 byte_end 用 u64 打包，calc() 直接拆出来用：

// 改动后：word_id 从边上直接拿
for (byte_end, word_id) in dag.iter_edges(byte_start) {
    let freq = if word_id != NO_MATCH {
        self.records[word_id as usize].freq
    } else {
        1
    };
}

编码方式：byte_end + 1 放 u64 高 32 位，word_id（i32）放低 32 位，全零表示没有更多边了。

还有两个小改动。chars().count() == 1 换成 chars().nth(1).is_none()，前者遍历整个字符串计数，后者看一眼第二个字符就够了。HMM 路径里冗余的 re_han.is_match(block) 换成 state.is_matched()，split iterator 已经分类过了，没必要再跑一次正则。

干掉正则引擎 (PR #144)

PR #141 之后跑了一次 profiler，正则引擎还占 29% 的 CPU 时间。

这个数字让我有点意外。看看这几个正则在干嘛：RE_HAN_DEFAULT 匹配 CJK 字符加 ASCII 字母数字加几个标点，RE_SKIP_DEFAULT 匹配空白。没有回溯，没有量词嵌套，没有捕获组，就是判断一个码位在不在某个范围里。

regex crate 的 DFA 实现确实快，但通用引擎对这种需求来说还是太重了。换成 matches! 宏做字符分类，编译后就是几个范围比较，LLVM 直接变跳转表：

fn is_cjk(c: char) -> bool {
    matches!(c,
        '\u{3400}'..='\u{4DBF}'
        | '\u{4E00}'..='\u{9FFF}'
        | '\u{F900}'..='\u{FAFF}'
        | '\u{20000}'..='\u{2A6DF}'
        | '\u{2A700}'..='\u{2B73F}'
        | '\u{2B740}'..='\u{2B81F}'
        | '\u{2B820}'..='\u{2CEAF}'
        | '\u{2CEB0}'..='\u{2EBEF}'
        | '\u{2F800}'..='\u{2FA1F}'
    )
}

fn is_han_default(c: char) -> bool {
    is_cjk(c) || c.is_ascii_alphanumeric() || matches!(c, '+' | '#' | '&' | '.' | '_' | '%' | '-')
}

配套写了个 SplitByCharacterClass 泛型迭代器，接受 Fn(char) -> bool 分类函数，把文本分成 matched/unmatched 连续段，大概 30 行代码：

pub(crate) struct SplitByCharacterClass<'t, F> {
    text: &'t str,
    pos: usize,
    classify: F,
}

逐字符扫描，classify(c) 为 true 就延伸 matched 段，为 false 就延伸 unmatched 段。类别翻转时产出当前切片，开始新段。

RE_SKIP_DEFAULT 的空白匹配更简单，tokenize 里直接逐字符遍历，\r\n 配对处理，连 splitter 都不需要。

唯一保留的正则是 hmm.rs 里的 RE_SKIP，匹配 [a-zA-Z0-9]+(?:.\d+)?%?，有量词和可选组，用正则合理。

结果：所有 cut 模式快了 25-29%。no_hmm 从 1.45 µs 到 1.03 µs。正则引擎从占 29% CPU 变成几乎没有。

char 键 emit 概率，预计算 ln() (PR #145)

干掉正则之后又看了一次 profiler。

Viterbi 内层循环还是热的。每个字符在 4 个 HMM 状态下查一次发射概率，用 phf::Map。key 类型是 &str，但每个 key 其实就一个汉字，UTF-8 编码 3 字节。每次查找在哈希 3 字节的字符串、逐字节比较。换成 phf::Map<char, f64> 后变成哈希一个 u32。

需要改生成 HMM 数据的 proc macro：

// 改动前
pub static EMIT_PROB_0: phf::Map<&'static str, f64> = ...;

// 改动后
pub static EMIT_PROB_0: phf::Map<char, f64> = ...;

另一个热点是 calc() 对每条边算 (freq as f64).ln()。词频在字典加载后就不变了，每次分词重算没有意义。Record 结构体加个 log_freq 字段，构造时算好：

struct Record {
    freq: usize,
    log_freq: f64,   // (freq as f64).ln()，只算一次
    tag: Box<str>,
}

calc() 直接读 self.records[word_id].log_freq。频率 1 的回退值也预计算成 log1 = 0.0 - logtotal，提到循环外面。

Viterbi 循环本身也简化了。之前用 Peekable 迭代器手动管理字节偏移来取每个字符的 substring：

// 改动前
let mut curr = char_offsets.iter().copied().peekable();
let x1 = curr.next().unwrap();
let x2 = *curr.peek().unwrap();
// ... 手动 peek/next

现在收集 char_indices()，下标循环就行了：

// 改动后
let chars: Vec<(usize, char)> = sentence.char_indices().collect();
for t in 1..C {
    let ch = chars[t].1;
    let em_prob = params.emit_prob(*y as usize, ch);
    // ...
}

Vec 分配通过 HmmContext 复用摊还，很少真的走 allocator。

合起来 9-13%。with_hmm 从 1.49 µs 到 1.32 µs。

回头看

最重要的两个改动："连续整数 key 别用 HashMap"和"字符范围检查别用正则引擎"。这两个贡献了总提升的大部分，其余的单个不大，但叠加起来也有 10% 多。

每次都是 profiler 告诉我该看哪。不看火焰图我很容易猜错。正则引擎占 29% 这件事就很意外，那些 pattern 看着简单，regex crate 优化也到位了，但通用方案在需求简单到只需要几个范围比较的时候，还是比不过专用代码。

Rust 里写性能敏感代码，我对分配这件事反复犯同一个错误：热路径里 Vec::new() 在循环里面就是坏味道。挪出去，复用，resize。borrow checker 会烦你，但分配开销是实打实省下来的。

感谢 Runji Wang 和 aniaan 贡献了最初几个 PR。