> 拼音查询的 110 行 trie 代码、词频 boost、用户词典原子持久化。
本篇是 《TUX IM 开发日志:从零到 0.1》 系列第 4/8 篇。
第 4 篇 · 拼音模式:trie、词频、用户词典
第 4 篇 · 拼音模式:trie、词频、用户词典
这一篇开始技术深潜第一站。拼音查询是整个输入法最热路径——每个按键都要查 trie 几次。我会拆三件事:trie 数据结构本身、词频学习机制、用户词典的原子持久化。
trie:110 行撑起整个拼音查询
先看代码(tux_im/input/lexicon/_trie.py):
Trie 核心数据结构
@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()),然后对每个音节:
候选生成(简化)
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 实现
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_first 和 select_candidate 都会调,ae1084b 加的。每次选候选 +10 freq。
> 坑:boost 和 insert 不要混用。如果同一个 (code, word) 既有内置词典的 freq=9999,又被 boost 到 100000,下次启动词典加载时 freq 会被覆盖回 9999——除非走 add_user_word 路径而不是 boost。d9c5f29 之后 commit_first/select_candidate 改成调 add_user_word,问题就消失了。
用户词典的原子持久化
用户学习过的词要保存到 ~/.local/share/tux-im/user_words.tsv。问题是:写到一半进程被杀怎么办?
三个关键设计
- dirty flag + 防抖——每次 add_user_word 标记 dirty,5s 后才 flush。不每次写。
- 原子写——写到
.tmp再os.replace(),POSIX rename 是原子的。 - atexit 注册——进程正常退出时强制 flush 一次。
原子 flush
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 格式)一起加载。优先级:
加载顺序
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):所有三个方法(candidates、select、commit)都用完整音节查 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 同步这两个有意思的设计。
评论