msime_engine/lattice/
octagram.rs

1//! Rime octagram（八股文）`.gram` 语法模型的只读查找端。
2//!
3//! 出处有三处，许可证各不相同：
4//!
5//! - 文件格式与查询算法来自 lotem/librime-octagram，固定于提交 `57d18b9f58e5284bd891d559f6bdd16cf60341e9`，Copyright (c) 2014, RIME Developers，BSD-3-Clause，完整许可见 `resources/licenses/librime-octagram-BSD-3-Clause.txt`。
6//! - 代码逐函数翻译自 metasequoiaime/msime-windows 的 `engine/ngram/octagram/octagram_gram.cpp` 与 `octagram_gram.h`，作者 jsfaint，以 GPL-3.0 发布，作者已同意随引擎以 MPL-2.0 再许可（见 `crates/engine/README.md` 的「许可证」一节）；部分注释沿用了那里的措辞。
7//! - 双数组 Trie 的读取（`traverse` 与 `DoubleArrayUnit` 位段）按 darts-clone 0.32 的数据布局重新实现，darts-clone 为 Copyright (c) 2008-2014, Susumu Yata，BSD-2-Clause，许可证全文附在同一份 `resources/licenses/librime-octagram-BSD-3-Clause.txt` 的末尾。
8//!
9//! 文件是「Metadata + Darts 双数组 Trie 镜像」：键为「上下文尾部字符 + 词首字符」经 `encode` 压缩的字节串，值为 `int(ln(共现频次) * 10000)`。它不是概率语言模型：`query` 给出的是搭配强度信号——命中取 `max(值/10000 + 搭配惩罚)`，全部未命中返回 `non_collocation_penalty`。词表零耦合：我们词库里没有的词只是拿不到加成，不会报错。
10//!
11//! 头部布局（rime `MappedFile` 上的 Metadata，偏移 0）：`format` char\[32\]、uint32 checksum、uint32 `double_array_size`（单位数，非字节）、int32 `OffsetPtr<char>`（相对自身字段的偏移）。双数组单元 4 字节。
12//!
13//! 装载沿用 `ngram.rs` 的 mmap 约定：`.gram` 经资源包以原子重命名安装
14//! （`.incoming` 改名就位，绝不就地改写），映射随对象存活；万象模型的镜像接近
15//! 400 MB，整读进内存不可取，RSS 只随命中的页增长。
16
17use std::num::NonZeroUsize;
18use std::path::{Path, PathBuf};
19use std::sync::{Arc, Mutex, OnceLock, PoisonError};
20
21use lru::LruCache;
22use memmap2::{Mmap, MmapOptions};
23
24/// `query` 返回自然对数域的搭配强度；词格主分是 log10 口径，装配方乘
25/// `f64::consts::LOG10_E`（= 1/ln(10)）换域，weight 才是纯标定系数。
26/// 打分参数。默认 = 万象现行（`wanxiang.schema.yaml` 的 6/2/-14/-6/-100/-18），
27/// 与旧版出厂模板 `config.default.toml` 的 `[association]` 键一致。octagram
28/// 自身默认 `collocation_max_length=4`、`collocation_min_length=3`、
29/// `collocation_penalty = non_collocation_penalty = -12`，我们不采用。
30#[derive(Debug, Clone, Copy, PartialEq)]
31pub struct GramConfig {
32    /// 上下文与词首各取 `collocation_max_length - 1` 个码点，上限 8。
33    pub collocation_max_length: i32,
34    /// 搭配长度（上下文现存码点数 + 命中键码点数）低于它时吃更狠的 weak 惩罚。
35    pub collocation_min_length: i32,
36    pub collocation_penalty: f64,
37    /// 全部窗口未命中、上下文为空或模型无效时的常数分。
38    pub non_collocation_penalty: f64,
39    pub weak_collocation_penalty: f64,
40    /// 句尾词的「词+$」rear 项惩罚。
41    pub rear_penalty: f64,
42}
43
44impl Default for GramConfig {
45    fn default() -> Self {
46        Self {
47            collocation_max_length: 6,
48            collocation_min_length: 2,
49            collocation_penalty: -14.0,
50            non_collocation_penalty: -6.0,
51            weak_collocation_penalty: -100.0,
52            rear_penalty: -18.0,
53        }
54    }
55}
56
57/// `gram_encoding.cc` / `gram_db.cc` 的常量。
58pub const MAX_ENCODED_UNICODE: usize = 8;
59const MAX_RESULTS: usize = 8;
60const VALUE_SCALE: f64 = 10000.0;
61
62/// 双数组单元数上限：万象镜像约 400 MB（一亿单元），放宽到它的四倍挡恶意计数。
63const MAX_UNITS: usize = 400_000_000;
64
65const FORMAT_MAGIC: &[u8] = b"Rime::Grammar/";
66const HEADER_BYTES: usize = 44;
67const FORMAT_OFFSET: usize = 0;
68const ARRAY_SIZE_OFFSET: usize = 36;
69const OFFSET_PTR_OFFSET: usize = 40;
70const UNIT_BYTES: usize = 4;
71
72/// 一条 `.gram` 镜像：mmap 只读视图 + 双数组单元的字节区间。
73pub struct GramDb {
74    bytes: Mmap,
75    /// 双数组镜像起点在映射里的字节偏移。
76    array_at: usize,
77    /// 单元数（4 字节一单元）。
78    array_units: usize,
79}
80
81/// `common_prefix_search` 的一条命中：值与命中键的字节长度（不是码点数）。
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub struct GramMatch {
84    pub value: i32,
85    pub length: usize,
86}
87
88static SHARED: OnceLock<Mutex<LruCache<PathBuf, Option<Arc<GramDb>>>>> = OnceLock::new();
89const SHARED_CACHE_CAPACITY: usize = 4;
90
91impl GramDb {
92    /// 文件缺失、过短、魔数错误或数组越界时返回 `None`；调用方按无模型运行，
93    /// 与缺少 `msime-bigram.bin` 时的处理一致。
94    #[allow(unsafe_code)]
95    pub fn load(path: &Path) -> Option<GramDb> {
96        if !std::fs::symlink_metadata(path).ok()?.file_type().is_file() {
97            return None;
98        }
99        let mut file = crate::paths::open_file_no_follow(path).ok()?;
100        let size = file.metadata().ok()?.len() as usize;
101        if size < HEADER_BYTES {
102            return None;
103        }
104        // 头 44 字节先读出来校验，校验过了才把整文件映射上：魔数不对的文件连映射都不建。
105        use std::io::Read;
106        let mut header = [0u8; HEADER_BYTES];
107        file.read_exact(&mut header).ok()?;
108        if &header[FORMAT_OFFSET..FORMAT_OFFSET + FORMAT_MAGIC.len()] != FORMAT_MAGIC {
109            return None;
110        }
111        let array_units = u32::from_le_bytes(
112            header[ARRAY_SIZE_OFFSET..ARRAY_SIZE_OFFSET + 4]
113                .try_into()
114                .expect("four size bytes"),
115        ) as usize;
116        let offset = i32::from_le_bytes(
117            header[OFFSET_PTR_OFFSET..OFFSET_PTR_OFFSET + 4]
118                .try_into()
119                .expect("four offset bytes"),
120        );
121        if offset < 0 || array_units > MAX_UNITS {
122            return None;
123        }
124        let array_at = OFFSET_PTR_OFFSET + offset as usize;
125        if array_at > size || array_units > (size - array_at) / UNIT_BYTES {
126            return None;
127        }
128        // SAFETY: 与 `NgramTable::load` 同一份映射契约——`.gram` 由资源包以
129        // `.incoming` 改名安装，文件字节在映射存活期间不变；镜像边界上面已校验。
130        let bytes = unsafe { MmapOptions::new().map(&file) }.ok()?;
131        Some(GramDb {
132            bytes,
133            array_at,
134            array_units,
135        })
136    }
137
138    /// 按路径共享模型实例；缺失或非法的文件同样被记住（`None`），不会在每次
139    /// `set_sentence_association` 后重复探测。与 `NgramTable::shared` 同一范式。
140    pub fn shared(path: &Path) -> Option<Arc<GramDb>> {
141        let mut loaded = SHARED
142            .get_or_init(|| {
143                Mutex::new(LruCache::new(
144                    NonZeroUsize::new(SHARED_CACHE_CAPACITY).unwrap(),
145                ))
146            })
147            .lock()
148            .unwrap_or_else(PoisonError::into_inner);
149        if let Some(cached) = loaded.get(path) {
150            return cached.clone();
151        }
152        let db = GramDb::load(path).map(Arc::new);
153        loaded.put(path.to_path_buf(), db.clone());
154        db
155    }
156
157    /// 摘掉指定路径的缓存引用（设置侧删除模型包后调用）；仍被会话持有的实例
158    /// 由 `Arc` 兜底，最后一个引用松开才真正 unmap，被删的 `.gram` 随之可物理删除。
159    /// 资源包卸载通过 Host API 调用；活跃会话的引用仍保留到会话重建。
160    pub fn evict_shared(path: &Path) {
161        if let Some(cache) = SHARED.get() {
162            cache
163                .lock()
164                .unwrap_or_else(PoisonError::into_inner)
165                .pop(path);
166        }
167    }
168
169    #[cfg(test)]
170    fn unit_count(&self) -> usize {
171        self.array_units
172    }
173
174    /// 越界返回 `None` 而不是 panic：损坏的镜像最多损失一次搭配命中。
175    fn unit(&self, index: usize) -> Option<u32> {
176        if index >= self.array_units {
177            return None;
178        }
179        let at = self.array_at + index * UNIT_BYTES;
180        Some(u32::from_le_bytes(
181            self.bytes[at..at + UNIT_BYTES]
182                .try_into()
183                .expect("four unit bytes"),
184        ))
185    }
186
187    // ---- darts 双数组（darts-clone 0.32）的只读部分 --------------------------
188
189    /// `traverse`：从 `node_pos` 出发按 `key` 逐字节走转移，走完返回落点。
190    /// 中途标签不符或越界都返回 `None`。
191    fn traverse(&self, key: &[u8], mut node_pos: usize) -> Option<usize> {
192        let mut id = node_pos;
193        let mut unit = self.unit(id)?;
194        for &byte in key {
195            id ^= (offset_of(unit) as usize) ^ usize::from(byte);
196            unit = self.unit(id)?;
197            if label_of(unit) != u32::from(byte) {
198                return None;
199            }
200            node_pos = id;
201        }
202        Some(node_pos)
203    }
204
205    /// `commonPrefixSearch`：从 `node_pos` 出发沿 `key` 走，每个节点若有 `\0`
206    /// 叶孩子就记录一条命中（值 = 叶单元的 value，长度 = 已匹配键的字节数）。
207    fn common_prefix_search(&self, key: &[u8], mut node_pos: usize, results: &mut Vec<GramMatch>) {
208        let mut unit = match self.unit(node_pos) {
209            Some(unit) => unit,
210            None => return,
211        };
212        node_pos ^= offset_of(unit) as usize;
213        for (index, &byte) in key.iter().enumerate() {
214            node_pos ^= usize::from(byte);
215            let Some(next) = self.unit(node_pos) else {
216                return;
217            };
218            unit = next;
219            if label_of(unit) != u32::from(byte) {
220                return;
221            }
222            node_pos ^= offset_of(unit) as usize;
223            if has_leaf(unit) && results.len() < MAX_RESULTS {
224                if let Some(leaf) = self.unit(node_pos) {
225                    results.push(GramMatch {
226                        value: leaf as i32 & 0x7FFF_FFFF,
227                        length: index + 1,
228                    });
229                }
230            }
231        }
232    }
233
234    /// `GramDb::Lookup`：先 `traverse` 整段上下文，再从落点做词首的公共前缀查找。
235    fn lookup(&self, context_encoded: &[u8], word_encoded: &[u8], results: &mut Vec<GramMatch>) {
236        if context_encoded.is_empty() {
237            return;
238        }
239        let Some(node_pos) = self.traverse(context_encoded, 0) else {
240            return;
241        };
242        self.common_prefix_search(word_encoded, node_pos, results);
243    }
244
245    /// `Octagram::Query`。`context`、`word` 都是原始 UTF-8；`is_rear` 对应句尾
246    /// 词的「词+$」rear 项。返回值是自然对数域的搭配强度，未命中为
247    /// `config.non_collocation_penalty`。
248    pub fn query(&self, context: &str, word: &str, is_rear: bool, config: &GramConfig) -> f64 {
249        let mut result = config.non_collocation_penalty;
250        if context.is_empty() {
251            return result;
252        }
253        // 上下文取尾部 collocation_max_length-1 个码点，词取头部同样数量，
254        // 再各自 encode。collocation_len = 上下文现存码点数 + 命中键的码点数，
255        // 决定吃 collocation 惩罚还是更狠的 weak 惩罚。
256        let n = MAX_ENCODED_UNICODE
257            .min(config.collocation_max_length.saturating_sub(1).max(0) as usize);
258        if n < 1 {
259            return result;
260        }
261        let context_tail = last_codepoints(context, n);
262        let mut remaining_windows = codepoint_count(context_tail) as i32;
263        let context_query = encode(context_tail);
264        let word_head = first_codepoints(word, n);
265        let word_head_len = codepoint_count(word_head) as i32;
266        let word_query = encode(word_head);
267
268        let mut context_ptr = 0usize;
269        let mut matches = Vec::with_capacity(MAX_RESULTS);
270        while remaining_windows > 0 && context_ptr < context_query.len() {
271            let at_window_start = context_ptr == 0;
272            matches.clear();
273            self.lookup(&context_query[context_ptr..], &word_query, &mut matches);
274            for m in &matches {
275                let scaled = scale_value(m.value);
276                let match_len = unicode_length(&word_query, m.length) as i32;
277                let collocation_len = remaining_windows + match_len;
278                let whole_query = at_window_start && m.length == word_query.len();
279                let penalty = if collocation_len >= config.collocation_min_length || whole_query {
280                    config.collocation_penalty
281                } else {
282                    config.weak_collocation_penalty
283                };
284                result = result.max(scaled + penalty);
285            }
286            remaining_windows -= 1;
287            context_ptr = next_unicode(&context_query, context_ptr);
288        }
289
290        if is_rear && word_head_len == codepoint_count(word) as i32 {
291            // 句尾标记只在「词未被头部截断」时查：截过的词查 $ 无意义。
292            matches.clear();
293            self.lookup(&word_query, b"$", &mut matches);
294            if let Some(first) = matches.first() {
295                result = result.max(scale_value(first.value) + config.rear_penalty);
296            }
297        }
298        result
299    }
300}
301
302/// `scale_value`：负值不可能是合法频次分，一律按 -1 处理。
303fn scale_value(value: i32) -> f64 {
304    if value >= 0 {
305        value as f64 / VALUE_SCALE
306    } else {
307        -1.0
308    }
309}
310
311// ---- 双数组单元的位段（darts-clone 的 DoubleArrayUnit）------------------------
312
313fn label_of(unit: u32) -> u32 {
314    unit & 0x8000_00FF
315}
316
317fn has_leaf(unit: u32) -> bool {
318    (unit >> 8) & 1 == 1
319}
320
321fn offset_of(unit: u32) -> u32 {
322    (unit >> 10) << ((unit & 0x200) >> 6)
323}
324
325// ---- gram_encoding.cc 的移植 -------------------------------------------------
326
327/// 读一个原始 UTF-8 码点并推进 `i`。续字节缺位或非法序列按 Latin-1 单字节
328/// 兜底：搭配查询遇到乱码最多损失一次命中，不值得为它抛错。
329fn next_codepoint(text: &str, i: &mut usize) -> u32 {
330    let bytes = text.as_bytes();
331    let lead = bytes[*i];
332    *i += 1;
333    // 续字节：只在确实存在且形态合法时消费，值按需拼接。
334    let take = |i: &mut usize| -> u32 {
335        if *i < bytes.len() && (bytes[*i] & 0xC0) == 0x80 {
336            let b = bytes[*i];
337            *i += 1;
338            u32::from(b & 0x3F)
339        } else {
340            0
341        }
342    };
343    if lead < 0x80 {
344        return u32::from(lead);
345    }
346    if (lead & 0xE0) == 0xC0 {
347        return (u32::from(lead & 0x1F) << 6) | take(i);
348    }
349    if (lead & 0xF0) == 0xE0 {
350        let high = take(i);
351        let low = take(i); // 两次调用必须先后求值，不能写进同一个表达式
352        return (u32::from(lead & 0x0F) << 12) | (high << 6) | low;
353    }
354    if (lead & 0xF8) == 0xF0 {
355        let b1 = take(i);
356        let b2 = take(i);
357        let b3 = take(i);
358        return (u32::from(lead & 0x07) << 18) | (b1 << 12) | (b2 << 6) | b3;
359    }
360    u32::from(lead)
361}
362
363pub fn encode(text: &str) -> Vec<u8> {
364    let mut out = Vec::with_capacity(text.len());
365    let mut i = 0;
366    while i < text.len() {
367        let u = next_codepoint(text, &mut i);
368        if u < 0x80 {
369            // 编码串里不允许出现 NUL（历史格式按 C 字符串用），0 被转成 0xE0 占位。
370            out.push(if u == 0 { 0xE0 } else { u as u8 });
371        } else if (0x4000..0xA000).contains(&u) {
372            // CJK 常用区压成 2 字节：首字节平移进 0x4080-0xA07F，避开 ASCII 与
373            // 0xE0 变长前缀。
374            if (u & 0xFF) == 0 {
375                out.push(0xE1);
376                out.push(((u >> 8) + 0x40) as u8);
377            } else {
378                out.push(((u >> 8) + 0x40) as u8);
379                out.push((u & 0xFF) as u8);
380            }
381        } else {
382            let mut u = u;
383            let mut bits = 32i32;
384            while bits > 0 && (u & 0xFE00_0000) == 0 {
385                bits -= 7;
386                u <<= 7;
387            }
388            let mut bytes_to_encode = (bits + 6) / 7;
389            out.push(0xE0 | bytes_to_encode as u8);
390            while bytes_to_encode > 0 {
391                bytes_to_encode -= 1;
392                out.push((((u >> 25) & 0x7F) | 0x80) as u8);
393                u <<= 7;
394            }
395        }
396    }
397    out
398}
399
400/// 编码串上的游标推进（`advance`）：一个编码码点占 1 字节（ASCII）、2 字节
401/// （CJK 常用区）或 1+首字节低 4 位（扩展区变长）。越界时停在末尾。
402fn next_unicode(encoded: &[u8], pos: usize) -> usize {
403    let Some(&first) = encoded.get(pos) else {
404        return pos;
405    };
406    let step = if (first & 0x80) == 0 {
407        1
408    } else if (first & 0xF0) == 0xE0 {
409        usize::from(first & 0x0F) + 1
410    } else {
411        2
412    };
413    (pos + step).min(encoded.len())
414}
415
416/// 编码串前 `byte_length` 字节里的码点数（`unicode_length`）。
417fn unicode_length(encoded: &[u8], byte_length: usize) -> usize {
418    let mut len = 0;
419    let mut i = 0;
420    while i < byte_length && i < encoded.len() {
421        i = next_unicode(encoded, i);
422        len += 1;
423    }
424    len
425}
426
427/// 原始 UTF-8 串的码点数与首/尾 n 个码点。
428pub fn codepoint_count(text: &str) -> usize {
429    let mut n = 0;
430    let mut i = 0;
431    while i < text.len() {
432        next_codepoint(text, &mut i);
433        n += 1;
434    }
435    n
436}
437
438pub fn first_codepoints(text: &str, n: usize) -> &str {
439    let mut seen = 0;
440    let mut i = 0;
441    while i < text.len() && seen < n {
442        next_codepoint(text, &mut i);
443        seen += 1;
444    }
445    &text[..i]
446}
447
448pub fn last_codepoints(text: &str, n: usize) -> &str {
449    let total = codepoint_count(text);
450    if total <= n {
451        return text;
452    }
453    let mut skip = total - n;
454    let mut i = 0;
455    while i < text.len() && skip > 0 {
456        next_codepoint(text, &mut i);
457        skip -= 1;
458    }
459    &text[i..]
460}
461
462#[cfg(test)]
463pub(in crate::lattice) mod fixture {
464    // ---- 测试侧的双数组构造器 ------------------------------------------------
465    //
466    // 只为 fixture 存在：把「编码键 -> 值」排成一棵字节前缀树，再按 darts-clone
467    // 的布局规则落成单元数组。每个节点试探最小可用 offset（<256，免移位标志位），
468    // 子槽 = 节点下标 ^ offset ^ 子标签；键的末尾以 \0 叶单元存值，父单元置
469    // has_leaf 位。布局合法性的判据是读取端能逐项查回。
470
471    use std::collections::BTreeMap;
472    use std::path::{Path, PathBuf};
473
474    use super::{FORMAT_MAGIC, HEADER_BYTES, UNIT_BYTES};
475
476    #[derive(Default)]
477    pub struct FixtureNode {
478        children: BTreeMap<u8, FixtureNode>,
479        value: Option<i32>,
480    }
481
482    impl FixtureNode {
483        pub fn insert(&mut self, key: &[u8], value: i32) {
484            let mut node = self;
485            for &byte in key {
486                node = node.children.entry(byte).or_default();
487            }
488            node.value = Some(value);
489        }
490    }
491
492    /// 把 `entries` 的（编码键, 值）写成一份合法的 `.gram` 文件字节。
493    pub fn build_gram(entries: &[(Vec<u8>, i32)]) -> Vec<u8> {
494        let mut root = FixtureNode::default();
495        for (key, value) in entries {
496            assert!(!key.is_empty());
497            root.insert(key, *value);
498        }
499        // BFS 布局：每个节点带自己的下标、入边标签与挑好的 offset。
500        struct Layout {
501            node: FixtureNode,
502            index: usize,
503            label: u32,
504            offset: u32,
505        }
506        let mut nodes = vec![Layout {
507            node: root,
508            index: 0,
509            label: 0,
510            offset: 0,
511        }];
512        let mut units: BTreeMap<usize, u32> = BTreeMap::new();
513        let mut used: std::collections::HashSet<usize> = [0].into_iter().collect();
514        let mut frontier = 0;
515        while frontier < nodes.len() {
516            let node_index = nodes[frontier].index;
517            let (labels, has_leaf_child) = {
518                let node = &nodes[frontier].node;
519                (
520                    node.children.keys().copied().collect::<Vec<u8>>(),
521                    node.value.is_some(),
522                )
523            };
524            // 需要的槽位：每个子标签一个；有值时另加 label-0 的叶槽。
525            let mut needed: Vec<u8> = labels.clone();
526            if has_leaf_child {
527                needed.push(0);
528            }
529            let mut offset = 1u32;
530            let slots = loop {
531                assert!(offset < 256, "fixture trie too dense for a small offset");
532                let slots: Vec<usize> = needed
533                    .iter()
534                    .map(|&label| node_index ^ offset as usize ^ usize::from(label))
535                    .collect();
536                let distinct = slots.iter().collect::<std::collections::HashSet<_>>();
537                if distinct.len() == slots.len() && slots.iter().all(|slot| !used.contains(slot)) {
538                    break slots;
539                }
540                offset += 1;
541            };
542            nodes[frontier].offset = offset;
543            // 叶槽（label 0）与子槽一并占用；叶单元 = MSB | value。
544            for (&label, &slot) in needed.iter().zip(&slots) {
545                used.insert(slot);
546                if label == 0 {
547                    units.insert(
548                        slot,
549                        0x8000_0000 | nodes[frontier].node.value.unwrap() as u32,
550                    );
551                }
552            }
553            for (&label, &slot) in labels.iter().zip(&slots[..labels.len()]) {
554                let child = std::mem::take(nodes[frontier].node.children.get_mut(&label).unwrap());
555                nodes.push(Layout {
556                    node: child,
557                    index: slot,
558                    label: u32::from(label),
559                    offset: 0,
560                });
561            }
562            frontier += 1;
563        }
564        // 写每个节点的单元：offset<<10 | has_leaf<<8 | 入边标签。
565        for node in &nodes {
566            let has_leaf = u32::from(node.node.value.is_some());
567            units.insert(
568                node.index,
569                (node.offset << 10) | (has_leaf << 8) | node.label,
570            );
571        }
572
573        let array_units = used.iter().copied().max().unwrap_or(0) + 1;
574        let mut bytes = Vec::with_capacity(HEADER_BYTES + array_units * UNIT_BYTES);
575        let mut format = [0u8; 32];
576        format[..FORMAT_MAGIC.len()].copy_from_slice(FORMAT_MAGIC);
577        bytes.extend_from_slice(&format);
578        bytes.extend_from_slice(&0u32.to_le_bytes()); // checksum 沿袭上游不校验
579        bytes.extend_from_slice(&(array_units as u32).to_le_bytes());
580        bytes.extend_from_slice(&4i32.to_le_bytes()); // OffsetPtr：相对字段偏移
581        for index in 0..array_units {
582            bytes.extend_from_slice(&units.get(&index).copied().unwrap_or(0).to_le_bytes());
583        }
584        bytes
585    }
586
587    pub fn gram_file(dir: &Path, name: &str, entries: &[(Vec<u8>, i32)]) -> PathBuf {
588        let path = dir.join(name);
589        std::fs::write(&path, build_gram(entries)).unwrap();
590        path
591    }
592}
593
594#[cfg(test)]
595mod tests {
596    use super::fixture::*;
597    use super::*;
598
599    // ---- encode / 码点工具 ---------------------------------------------------
600
601    #[test]
602    fn encode_matches_gram_encoding() {
603        // ASCII 一字节。
604        assert_eq!(encode("abc"), b"abc");
605        // NUL 占位 0xE0。
606        assert_eq!(encode("\0"), vec![0xE0]);
607        // CJK 常用区两字节：中 U+4E2D -> [0x8E, 0x2D]，国 U+56FD -> [0x96, 0xFD]。
608        assert_eq!(encode("中国"), vec![0x8E, 0x2D, 0x96, 0xFD]);
609        // 低位为 0 的 CJK：一 U+4E00 -> [0xE1, 0x8E]。
610        assert_eq!(encode("一"), vec![0xE1, 0x8E]);
611        // 常用区之外（0xA000 边界、扩展区）走变长：é U+00E9 -> 0xE2 前缀 + 两个 7bit 组。
612        assert_eq!(encode("é"), vec![0xE2, 0x8E, 0xC8]);
613        // emoji U+1F600：3 个 7bit 组。
614        let emoji = encode("😀");
615        assert_eq!(emoji[0], 0xE0 | 3);
616        assert_eq!(emoji.len(), 4);
617    }
618
619    /// 真实 `.gram` 的人工验证探针：`GRAM_PROBE_PATH=<file>` 指向资源锁里的
620    /// 八股文紧凑模型，校验模型身份并对照独立读取器的分值。
621    #[test]
622    #[ignore = "需要锁定的八股文紧凑模型：GRAM_PROBE_PATH=<path> cargo test -p msime-engine --lib real_gram -- --ignored"]
623    fn real_gram_loads_and_queries() {
624        use sha2::{Digest, Sha256};
625
626        let path = std::env::var_os("GRAM_PROBE_PATH")
627            .expect("手动探针需要 GRAM_PROBE_PATH 指定锁定的八股文紧凑模型");
628        let db = GramDb::load(Path::new(&path)).expect("八股文紧凑模型能加载");
629        let lock: serde_json::Value = serde_json::from_str(include_str!(
630            "../../../../resources/grammar-models.lock.json"
631        ))
632        .expect("语法模型资源锁能解析");
633        let model = lock["models"]
634            .as_array()
635            .expect("资源锁包含模型列表")
636            .iter()
637            .find(|model| model["id"] == "baguwen-compact")
638            .expect("资源锁包含八股文紧凑模型");
639        let artifact = &model["artifacts"][0];
640        assert_eq!(
641            db.bytes.len() as u64,
642            artifact["size"].as_u64().expect("资源锁包含模型大小"),
643            "模型大小必须与锁定的八股文紧凑模型一致"
644        );
645        assert_eq!(
646            hex::encode(Sha256::digest(&db.bytes[..])),
647            artifact["sha256"].as_str().expect("资源锁包含模型 SHA-256"),
648            "模型 SHA-256 必须与锁定的八股文紧凑模型一致"
649        );
650
651        // 固定 darts-clone e40ce4627526985a7767444b6ed6893ab6ff8983 的独立读取结果。
652        // 模型分值以 1/10000 为步长，1e-9 容差只容纳浮点舍入。
653        let config = GramConfig::default();
654        for (context, word, expected_units) in [
655            ("一名", "医生", -40_443),
656            ("一名", "一生", -38_286),
657            ("北京", "大学", -3_926),
658            ("北京", "大雪", -26_763),
659        ] {
660            let expected = f64::from(expected_units) / 10_000.0;
661            let actual = db.query(context, word, true, &config);
662            eprintln!("{context}+{word}={actual:.4}，参考={expected:.4}");
663            assert!(
664                (actual - expected).abs() < 1e-9,
665                "{context}+{word} 的分值应为 {expected:?}，实际 {actual:?}"
666            );
667        }
668    }
669
670    #[test]
671    fn encoded_cursor_helpers() {
672        let enc = encode("中国ab一");
673        // 中[2] 国[2] a[1] b[1] 一[2]
674        assert_eq!(unicode_length(&enc, enc.len()), 5);
675        assert_eq!(next_unicode(&enc, 0), 2);
676        assert_eq!(next_unicode(&enc, 2), 4);
677        assert_eq!(next_unicode(&enc, 4), 5);
678        assert_eq!(next_unicode(&enc, 6), 8);
679        assert_eq!(next_unicode(&enc, 8), 8, "停在末尾不越界");
680        // 前缀字节内的码点数（前缀查找的 match.length 语义）。
681        assert_eq!(unicode_length(&enc, 4), 2);
682        assert_eq!(unicode_length(&enc, 3), 2, "截在半个码点处也算一个码点");
683    }
684
685    #[test]
686    fn raw_codepoint_helpers() {
687        assert_eq!(codepoint_count("中国人民"), 4);
688        assert_eq!(first_codepoints("中国人民", 2), "中国");
689        assert_eq!(last_codepoints("中国人民", 2), "人民");
690        assert_eq!(last_codepoints("中国", 8), "中国");
691        assert_eq!(first_codepoints("ab中", 2), "ab");
692    }
693
694    // ---- 装载与回退 -----------------------------------------------------------
695
696    #[test]
697    fn load_rejects_missing_and_broken_files() {
698        let dir = tempfile::tempdir().unwrap();
699        let d = dir.path();
700        assert!(GramDb::load(&d.join("absent.gram")).is_none(), "缺失文件");
701
702        let short = d.join("short.gram");
703        std::fs::write(&short, b"Rime::Gram").unwrap();
704        assert!(GramDb::load(&short).is_none(), "不足头部");
705
706        let magic = d.join("magic.gram");
707        std::fs::write(&magic, [b"Xime::Grammar/".as_slice(), &[0; 60]].concat()).unwrap();
708        assert!(GramDb::load(&magic).is_none(), "魔数不符");
709
710        // 魔数对但双数组镜像越界。
711        let oob = d.join("oob.gram");
712        let mut bytes = Vec::new();
713        let mut format = [0u8; 32];
714        format[..FORMAT_MAGIC.len()].copy_from_slice(FORMAT_MAGIC);
715        bytes.extend_from_slice(&format);
716        bytes.extend_from_slice(&0u32.to_le_bytes());
717        bytes.extend_from_slice(&1000u32.to_le_bytes());
718        bytes.extend_from_slice(&4i32.to_le_bytes());
719        bytes.extend_from_slice(&[0; 16]);
720        std::fs::write(&oob, &bytes).unwrap();
721        assert!(GramDb::load(&oob).is_none(), "镜像越界");
722
723        // 负偏移。
724        let neg = d.join("neg.gram");
725        let mut bytes = Vec::new();
726        bytes.extend_from_slice(&format);
727        bytes.extend_from_slice(&0u32.to_le_bytes());
728        bytes.extend_from_slice(&1u32.to_le_bytes());
729        bytes.extend_from_slice(&(-4i32).to_le_bytes());
730        bytes.extend_from_slice(&[0; 64]);
731        std::fs::write(&neg, &bytes).unwrap();
732        assert!(GramDb::load(&neg).is_none(), "负偏移");
733    }
734
735    // ---- 查询语义 -------------------------------------------------------------
736
737    /// 单条「上下文+词首」记录的模型：命中给 collocation_penalty 加成。
738    #[test]
739    fn query_scores_a_known_collocation() {
740        let dir = tempfile::tempdir().unwrap();
741        let config = GramConfig::default();
742        let key = [encode("中国"), encode("人民")].concat();
743        // ln(频次)*10000：取 12.0 -> 120000，比 collocation_penalty 高得多。
744        let path = gram_file(dir.path(), "m.gram", &[(key, 120000)]);
745        let db = GramDb::load(&path).expect("fixture loads");
746        assert!(db.unit_count() > 1);
747
748        let score = db.query("我爱中国", "人民", false, &config);
749        // 窗口从尾部 5 码点「爱中国」起，逐后缀试「国」命中：scaled 12 + (-14) = -2。
750        assert!(
751            (score - (-2.0)).abs() < 1e-6,
752            "尾部窗口命中后的分，实际 {score}"
753        );
754
755        let miss = db.query("我爱中国", "首都", false, &config);
756        assert_eq!(miss, config.non_collocation_penalty, "未命中给常数分");
757        assert_eq!(
758            db.query("", "人民", false, &config),
759            config.non_collocation_penalty
760        );
761    }
762
763    #[test]
764    fn query_uses_the_rear_entry_for_sentence_final_words() {
765        let dir = tempfile::tempdir().unwrap();
766        let config = GramConfig::default();
767        let path = gram_file(
768            dir.path(),
769            "rear.gram",
770            &[([encode("人民"), encode("$")].concat(), 160000)],
771        );
772        let db = GramDb::load(&path).expect("fixture loads");
773        // 词+$ 命中：scaled 16 + (-18) = -2，比 non_collocation -6 高。
774        assert!((db.query("无关上下文", "人民", true, &config) + 2.0).abs() < 1e-6);
775        // 非句尾不查 $。
776        assert_eq!(
777            db.query("无关上下文", "人民", false, &config),
778            config.non_collocation_penalty
779        );
780        // 词被头部截断时（词长 > collocation_max_length-1）不查 $。
781        let long_word = "人民万岁万岁万万岁";
782        assert_eq!(
783            db.query("无关", long_word, true, &config),
784            config.non_collocation_penalty
785        );
786    }
787
788    #[test]
789    fn query_bounds_extreme_configured_window_lengths() {
790        let dir = tempfile::tempdir().unwrap();
791        let path = gram_file(
792            dir.path(),
793            "window-limits.gram",
794            &[
795                ([encode("名"), encode("医")].concat(), 180000),
796                ([encode("一名"), encode("医")].concat(), 200000),
797                ([encode("医生"), encode("$")].concat(), 190000),
798            ],
799        );
800        let db = GramDb::load(&path).expect("合成模型能加载");
801        for (length, expected) in [
802            (i32::MIN, -6.0),
803            (i32::MIN + 1, -6.0),
804            (-5, -6.0),
805            (0, -6.0),
806            (1, -6.0),
807            (2, 4.0),
808            (3, 6.0),
809            (6, 6.0),
810            (9, 6.0),
811            (i32::MAX, 6.0),
812        ] {
813            let config = GramConfig {
814                collocation_max_length: length,
815                ..GramConfig::default()
816            };
817            for is_rear in [false, true] {
818                assert_eq!(
819                    db.query("一名", "医生", is_rear, &config),
820                    expected,
821                    "窗口长度={length}，句尾={is_rear}"
822                );
823            }
824        }
825    }
826
827    #[test]
828    fn query_takes_the_best_window() {
829        let dir = tempfile::tempdir().unwrap();
830        let config = GramConfig::default();
831        let entries = [
832            ([encode("国"), encode("人民")].concat(), 120000),
833            ([encode("我爱中国"), encode("人民")].concat(), 180000),
834        ];
835        let path = gram_file(dir.path(), "w.gram", &entries);
836        let db = GramDb::load(&path).expect("fixture loads");
837        // 首窗口「我爱中国」全键命中 18 + (-14) = 4，压过尾窗「国」的 12-14 = -2。
838        let score = db.query("我爱中国", "人民", false, &config);
839        assert!(
840            (score - 4.0).abs() < 1e-6,
841            "取所有窗口的最大值，实际 {score}"
842        );
843    }
844
845    #[test]
846    fn load_rejects_empty_array_outside_the_file() {
847        let dir = tempfile::tempdir().unwrap();
848        let mut bytes = vec![0u8; HEADER_BYTES];
849        bytes[..FORMAT_MAGIC.len()].copy_from_slice(FORMAT_MAGIC);
850        bytes[OFFSET_PTR_OFFSET..HEADER_BYTES].copy_from_slice(&1000i32.to_le_bytes());
851        let path = dir.path().join("outside.gram");
852        std::fs::write(&path, bytes).unwrap();
853        assert!(GramDb::load(&path).is_none());
854    }
855
856    #[test]
857    fn shared_caches_hits_and_misses() {
858        let dir = tempfile::tempdir().unwrap();
859        let missing = dir.path().join("missing.gram");
860        assert!(GramDb::shared(&missing).is_none());
861        let path = gram_file(
862            dir.path(),
863            "s.gram",
864            &[([encode("a"), encode("b")].concat(), 50000)],
865        );
866        let first = GramDb::shared(&path).expect("loads");
867        let second = GramDb::shared(&path).expect("same instance");
868        assert!(Arc::ptr_eq(&first, &second));
869        GramDb::evict_shared(&path);
870        let third = GramDb::shared(&path).expect("reloads after evict");
871        assert!(!Arc::ptr_eq(&first, &third));
872        GramDb::evict_shared(&path);
873    }
874}
msime_engine/lattice/octagram.rs

msime_engine/lattice/
octagram.rs
