> 拼音查询的 110 行 trie 代码、词频 boost、用户词典原子持久化。

本篇是 《TUX IM 开发日志:从零到 0.1》 系列第 4/8 篇。

第 4 篇 · 拼音模式:trie、词频、用户词典

第 4 篇 · 拼音模式:trie、词频、用户词典

这一篇开始技术深潜第一站。拼音查询是整个输入法最热路径——每个按键都要查 trie 几次。我会拆三件事:trie 数据结构本身、词频学习机制、用户词典的原子持久化。

trie:110 行撑起整个拼音查询

先看代码(tux_im/input/lexicon/_trie.py):

Trie 核心数据结构

Python
@dataclass
class _TrieNode:
    children: dict[str, _TrieNode] = field(default_factory=dict)
    entries: list[LexEntry] = field(default_factory=list)
    is_word: bool = False

class Trie:
    __slots__ = ("_root", "_size")

    def __init__(self) -> None:
        self._root = _TrieNode()
        self._size = 0

    def insert(self, code: str, word: str, freq: int = 0) -> None:
        code = code.lower().strip()
        if not code or not word:
            return
        node = self._root
        for ch in code:
            node = node.children.setdefault(ch, _TrieNode())
        if not node.is_word:
            self._size += 1
        node.entries.append(LexEntry(word=word, freq=freq, code=code))
        node.entries.sort(key=lambda e: (-e.freq, e.word))

整个文件 110 行,零 IO、零日志、零副作用。所有 14 个 trie 算法测试都不需要 mock 文件系统。

> 为什么纯数据结构是关键?因为 trie 是热路径——每次按键都查。如果 trie 里混了 IO,整个热路径就被拖慢;混了 logging,单元测试就会被打断;混了状态,测试就不可重复。纯 = 快 = 可测 = 可信赖

trie lookup 怎么做候选

对于 buffer nihao,先切分成 ['ni3', 'hao3']PinyinMode.segment()),然后对每个音节:

候选生成(简化)

Python
def candidates(self, limit=9):
    segments = self.segment(self.buffer)  # ['ni3', 'hao3']
    result = []
    for i, seg in enumerate(segments):
        entries = self._trie.lookup(seg)
        # 取 top-k,加 segment 索引
        for j, e in enumerate(entries[:limit]):
            result.append(Candidate(text=e.word, freq=e.freq, ...))
    return result[:limit]

实际上还会做笛卡尔积式的组合(ni3 的 top-5 × hao3 的 top-5),按总词频排序。这是 trie 数据结构之外的事,但 trie 给它提供了"前缀查 top-k"这个原子操作。

词频 boost:让"我"永远在第一

打完"我"按 1,下次再打"wo","我"还是第一个——因为 boost 了。

boost 实现

Python
def boost(self, code: str, word: str, delta: int = 10) -> None:
    code = code.lower().strip()
    node = self._root
    for ch in code:
        if ch not in node.children:
            return
        node = node.children[ch]
    if not node.is_word:
        return
    for e in node.entries:
        if e.word == word:
            e.freq += delta
            break
    node.entries.sort(key=lambda e: (-e.freq, e.word))

调用时机:commit_firstselect_candidate 都会调,ae1084b 加的。每次选候选 +10 freq。

> 坑:boost 和 insert 不要混用。如果同一个 (code, word) 既有内置词典的 freq=9999,又被 boost 到 100000,下次启动词典加载时 freq 会被覆盖回 9999——除非走 add_user_word 路径而不是 boostd9c5f29 之后 commit_first/select_candidate 改成调 add_user_word,问题就消失了。

用户词典的原子持久化

用户学习过的词要保存到 ~/.local/share/tux-im/user_words.tsv。问题是:写到一半进程被杀怎么办?

三个关键设计

  • dirty flag + 防抖——每次 add_user_word 标记 dirty,5s 后才 flush。不每次写。
  • 原子写——写到 .tmpos.replace(),POSIX rename 是原子的。
  • atexit 注册——进程正常退出时强制 flush 一次。
    原子 flush

Python
def _flush_now(self) -> None:
    if not self._dirty:
        return
    tmp = self._user_words_path.with_suffix('.tsv.tmp')
    tmp.write_text(self._serialize(), encoding='utf-8')
    os.replace(tmp, self._user_words_path)
    self._dirty = False

示意图

示意图

启动时的合并策略

用户词典(TSV)和系统词典(Rime 格式)一起加载。优先级:

加载顺序

Text
1. 加载 Rime 系统词典(只读)→ trie.insert(code, word, freq=系统freq)
2. 加载用户词典 TSV → 已存在的 entry: freq += 用户freq
                          新的 entry: trie.insert(code, word, freq=用户freq)

用户词典覆盖系统词典的 freq,但不删除系统词典的词。这就是为什么我打字"我"永远在第一——系统词典给"我" freq=9999,用户每次用 +10,最终达到 150000+,远远超过其他词。

那个差点让 0.1 发布出去的 bug

讲到 trie 必须提这个:tone digit 的去/留问题

问题:trie 里 key 是 ni3(带 tone),但有人在 candidates() 里写了 rstrip('12345'),查 ni 永远查不到带 tone 的 key。

修复(cddcb40):所有三个方法(candidatesselectcommit)都用完整音节查 trie。

教训:被这个 bug 抓到是因为 test_engine_flow.py 的 28 个测试。这个测试套件存在的意义就是——下次有人想"优化" trie 查询时,先看测试过没过

本篇小结

  • trie 是纯数据结构——零 IO、零日志、110 行。这是热路径的必需品。
  • boost 用于运行时调整add_user_word 用于持久化学习,两者路径不要混。
  • 原子写三件套:dirty flag + tmp+rename + atexit。
  • 启动合并顺序:系统词典 → 用户词典,用户 freq 覆盖系统 freq,不删词。
  • tone digit 是 trie key 的一部分,永远不要 strip。

> 下一篇五笔与混打:Wubi + Wbpy 两种哲学。讲讲五笔的"死键透传"和 wbpy 双 buffer 同步这两个有意思的设计。

最后修改: 2026年7月1日

作者

评论

发表评论

您的邮箱地址不会被公开。