lindera_dictionary/dictionary/prefix_dictionary.rs
1//! The prefix dictionary: surface forms mapped to their `WordEntry` records.
2//!
3//! The system dictionary ([`PrefixDictionary`]) is a char-wise double-array
4//! trie (built with crawdad) that is **walked in place** over the serialized
5//! bytes of `dict.trie` -- no deserialization, no owned node array. The trie
6//! maps a surface form to its key ordinal; `dict.valsidx` (a `u32` prefix sum)
7//! turns that ordinal into a run of records inside `dict.vals`.
8//!
9//! The user dictionary ([`UserPrefixDictionary`]) still uses a daachorse
10//! Aho-Corasick automaton, because prebuilt user-dictionary `.bin` files embed
11//! its serialized form inside an rkyv archive; changing it would invalidate
12//! every `.bin` in the wild for little gain (user dictionaries are small).
13
14use std::collections::BTreeMap;
15
16use byteorder::{ByteOrder, LittleEndian, WriteBytesExt};
17use daachorse::DoubleArrayAhoCorasick;
18use rkyv::rancor::{Fallible, Source};
19use rkyv::with::{ArchiveWith, DeserializeWith, SerializeWith};
20use rkyv::{Archive, Deserialize as RkyvDeserialize, Place, Serialize as RkyvSerialize};
21use serde::{Deserialize, Serialize};
22
23use crate::{LinderaResult, error::LinderaErrorKind, util::Data, viterbi::WordEntry};
24
25/// Match structure for common prefix iterator compatibility
26#[derive(Debug, Clone)]
27pub struct Match {
28 /// Which word matched.
29 pub word_idx: WordIdx,
30 /// Match length in characters (the number of `chars` consumed).
31 pub end_char: usize,
32}
33
34/// Identifies a word by its id within the dictionary it came from.
35#[derive(Debug, Clone, Copy)]
36pub struct WordIdx {
37 /// The word id.
38 pub word_id: u32,
39}
40
41impl WordIdx {
42 /// Wraps a raw word id.
43 ///
44 /// # Arguments
45 ///
46 /// * `word_id` - The word id to wrap.
47 ///
48 /// # Returns
49 ///
50 /// The wrapped id.
51 pub fn new(word_id: u32) -> Self {
52 Self { word_id }
53 }
54}
55
56/// Mask selecting the index/value bits of a crawdad node's `base`/`check`;
57/// the top bit is the leaf flag (`base`) / has-leaf flag (`check`).
58const OFFSET_MASK: u32 = 0x7fff_ffff;
59
60/// The code-mapper table value marking a character with no code assigned.
61const INVALID_CODE: u32 = u32::MAX;
62
63/// Byte offset where the code-mapper table starts inside `dict.trie`
64/// (after the `u32` table length).
65const TABLE_START: usize = 4;
66
67/// Byte length of one serialized trie node (`base: u32` + `check: u32`).
68const NODE_LEN_BYTES: usize = 8;
69
70/// The system prefix dictionary: a serialized crawdad trie walked in place.
71///
72/// Serialized trie layout (crawdad 0.3, `Trie::serialize_to_vec`):
73///
74/// ```text
75/// [table_len: u32][table: 4*table_len][alphabet_size: u32][node_len: u32][nodes: 8*node_len]
76/// ```
77///
78/// This layout is not a documented contract of the crawdad crate, so the
79/// dependency is pinned exactly and a round-trip test
80/// (`trie_view_matches_crawdad_search`) fails loudly if it ever changes.
81///
82/// All node and table reads go through checked slicing (`get`) and
83/// alignment-agnostic `u32` reads, so malformed bytes yield "no match" or a
84/// load-time error -- never a panic or undefined behaviour. That is what
85/// retired the O(n) validating pass and the `deserialize_unchecked` path the
86/// daachorse representation needed.
87#[derive(Clone)]
88pub struct PrefixDictionary {
89 /// The serialized trie (`dict.trie`), walked in place.
90 trie_data: Data,
91 /// Number of code-mapper table entries (indexable code points).
92 table_len: usize,
93 /// Byte offset of the node array inside `trie_data`.
94 nodes_start: usize,
95 /// Prefix sum over per-surface entry counts (`dict.valsidx`): `u32` LE
96 /// records, one per trie key plus a trailing sentinel, in units of
97 /// [`WordEntry::SERIALIZED_LEN`]-byte records inside `vals_data`.
98 vals_idx: Data,
99 /// The values file (`dict.vals`): `WordEntry` records back to back.
100 pub vals_data: Data,
101 /// Byte offsets into `words_data`, one `u32` per word id.
102 pub words_idx_data: Data,
103 /// Word detail records.
104 pub words_data: Data,
105}
106
107impl PrefixDictionary {
108 /// Serializes `word_entry_map` into `dict.trie` and `dict.valsidx` bytes.
109 ///
110 /// This is the build-time half of the dictionary: it is the only place
111 /// that runs crawdad itself. The caller is responsible for writing
112 /// `dict.vals` from the same map in the same iteration order (which a
113 /// `BTreeMap` fixes), so the prefix sum built here indexes it correctly.
114 ///
115 /// # Arguments
116 ///
117 /// * `word_entry_map` - Surface form to word entries, sorted and
118 /// deduplicated by the `BTreeMap`.
119 ///
120 /// # Returns
121 ///
122 /// `(trie_bytes, vals_idx_bytes)`, or an error if a surface is empty or
123 /// contains NUL (which crawdad reserves as its end marker), or if the
124 /// trie build fails.
125 pub fn serialize_trie(
126 word_entry_map: &BTreeMap<String, Vec<WordEntry>>,
127 ) -> LinderaResult<(Vec<u8>, Vec<u8>)> {
128 if word_entry_map.is_empty() {
129 // crawdad rejects an empty key set, so emit a minimal image the
130 // view treats as "matches nothing": an empty code table (every
131 // character maps to no code) and zero nodes.
132 let mut trie_bytes = Vec::with_capacity(12);
133 trie_bytes.extend_from_slice(&0u32.to_le_bytes()); // table_len
134 trie_bytes.extend_from_slice(&0u32.to_le_bytes()); // alphabet_size
135 trie_bytes.extend_from_slice(&0u32.to_le_bytes()); // node_len
136 return Ok((trie_bytes, 0u32.to_le_bytes().to_vec()));
137 }
138
139 let mut keys: Vec<&str> = Vec::with_capacity(word_entry_map.len());
140 let mut offsets: Vec<u32> = Vec::with_capacity(word_entry_map.len() + 1);
141 let mut acc: u32 = 0;
142
143 for (surface, entries) in word_entry_map {
144 // Skipping would silently desynchronize the prefix sum from the
145 // separately-written dict.vals, so a bad surface is an error.
146 // (An empty surface cannot get here anyway: the CSV parser drops
147 // empty fields before the map is built.)
148 if surface.is_empty() || surface.contains('\0') {
149 return Err(LinderaErrorKind::Build.with_error(anyhow::anyhow!(
150 "surface {surface:?} cannot be stored in the trie (empty or contains NUL)"
151 )));
152 }
153 keys.push(surface.as_str());
154 offsets.push(acc);
155 acc = acc.checked_add(entries.len() as u32).ok_or_else(|| {
156 LinderaErrorKind::Build.with_error(anyhow::anyhow!("entry offset overflowed u32"))
157 })?;
158 }
159 // Sentinel so the run for the last key is `idx[n-1]..idx[n]`.
160 offsets.push(acc);
161
162 let trie = crawdad::Trie::from_keys(keys.iter().copied()).map_err(|err| {
163 LinderaErrorKind::Build.with_error(anyhow::anyhow!("crawdad trie build failed: {err}"))
164 })?;
165 let trie_bytes = trie.serialize_to_vec();
166
167 let mut idx_bytes = Vec::with_capacity(offsets.len() * 4);
168 for offset in &offsets {
169 idx_bytes
170 .write_u32::<LittleEndian>(*offset)
171 .map_err(|err| {
172 LinderaErrorKind::Io
173 .with_error(anyhow::anyhow!(err))
174 .add_context("Failed to encode values index")
175 })?;
176 }
177
178 Ok((trie_bytes, idx_bytes))
179 }
180
181 /// Builds an in-memory dictionary from surface forms and entries.
182 ///
183 /// Serializes the values file alongside the trie so the result is
184 /// self-consistent. Used by the trainer and by tests; the production path
185 /// loads previously-built files via [`PrefixDictionary::load`] instead.
186 ///
187 /// # Arguments
188 ///
189 /// * `word_entry_map` - Surface form to word entries.
190 ///
191 /// # Returns
192 ///
193 /// A dictionary answering prefix queries over the map, with empty word
194 /// detail data.
195 pub fn from_word_entry_map(
196 word_entry_map: &BTreeMap<String, Vec<WordEntry>>,
197 ) -> LinderaResult<Self> {
198 let (trie_bytes, idx_bytes) = Self::serialize_trie(word_entry_map)?;
199
200 let mut vals_bytes = Vec::with_capacity(word_entry_map.len() * WordEntry::SERIALIZED_LEN);
201 for entries in word_entry_map.values() {
202 for entry in entries {
203 entry.serialize(&mut vals_bytes).map_err(|err| {
204 LinderaErrorKind::Serialize
205 .with_error(anyhow::anyhow!(err))
206 .add_context("Failed to serialize word entry")
207 })?;
208 }
209 }
210
211 Self::load(trie_bytes, idx_bytes, vals_bytes, Vec::new(), Vec::new())
212 }
213
214 /// Load a `PrefixDictionary` from raw binary data.
215 ///
216 /// Performs the O(1) structural check on the trie header: both length
217 /// headers must be consistent with the buffer, which is what stops a
218 /// crafted file from requesting an enormous allocation or walking out of
219 /// bounds. No O(n) node scan is needed -- every access during search is
220 /// bounds-checked, so a malformed node yields "no match" rather than a
221 /// panic.
222 ///
223 /// # Arguments
224 ///
225 /// * `trie_data` - Contents of `dict.trie`.
226 /// * `vals_idx` - Contents of `dict.valsidx`.
227 /// * `vals_data` - Contents of `dict.vals`.
228 /// * `words_idx_data` - Contents of `dict.wordsidx`.
229 /// * `words_data` - Contents of `dict.words`.
230 ///
231 /// # Returns
232 ///
233 /// A `PrefixDictionary`, or an error if the trie or index headers are
234 /// inconsistent with their buffers.
235 pub fn load(
236 trie_data: impl Into<Data>,
237 vals_idx: impl Into<Data>,
238 vals_data: impl Into<Data>,
239 words_idx_data: impl Into<Data>,
240 words_data: impl Into<Data>,
241 ) -> LinderaResult<PrefixDictionary> {
242 let trie_data = trie_data.into();
243 let vals_idx = vals_idx.into();
244
245 if trie_data.len() < TABLE_START {
246 return Err(LinderaErrorKind::Deserialize
247 .with_error(anyhow::anyhow!("dict.trie is too short for a trie header")));
248 }
249 let table_len = LittleEndian::read_u32(&trie_data[0..4]) as usize;
250 // TABLE_START (table_len) + table + 4 (alphabet_size) + 4 (node_len)
251 let nodes_start = table_len
252 .checked_mul(4)
253 .and_then(|table_bytes| table_bytes.checked_add(TABLE_START + 8))
254 .ok_or_else(implausible_size)?;
255 if nodes_start > trie_data.len() {
256 return Err(LinderaErrorKind::Deserialize.with_error(anyhow::anyhow!(
257 "dict.trie declares a {table_len}-entry code table that exceeds the file"
258 )));
259 }
260 let node_len = LittleEndian::read_u32(&trie_data[nodes_start - 4..nodes_start]) as usize;
261 let expected = node_len
262 .checked_mul(NODE_LEN_BYTES)
263 .and_then(|node_bytes| node_bytes.checked_add(nodes_start))
264 .ok_or_else(implausible_size)?;
265 if expected != trie_data.len() {
266 return Err(LinderaErrorKind::Deserialize.with_error(anyhow::anyhow!(
267 "dict.trie declares {node_len} nodes ({expected} bytes) but the file is {} bytes",
268 trie_data.len()
269 )));
270 }
271
272 if vals_idx.len() % 4 != 0 {
273 return Err(LinderaErrorKind::Deserialize.with_error(anyhow::anyhow!(
274 "dict.valsidx length {} is not a whole number of u32 records",
275 vals_idx.len()
276 )));
277 }
278
279 Ok(PrefixDictionary {
280 trie_data,
281 table_len,
282 nodes_start,
283 vals_idx,
284 vals_data: vals_data.into(),
285 words_idx_data: words_idx_data.into(),
286 words_data: words_data.into(),
287 })
288 }
289
290 /// Reads node `idx`'s `(base, check)` pair, raw (flag bits included).
291 ///
292 /// # Arguments
293 ///
294 /// * `idx` - Node index.
295 ///
296 /// # Returns
297 ///
298 /// The raw pair, or `None` when `idx` is out of range -- which for a
299 /// well-formed trie never happens, and for a malformed one safely ends
300 /// the search.
301 #[inline(always)]
302 fn node(&self, idx: u32) -> Option<(u32, u32)> {
303 let off = self.nodes_start + (idx as usize) * NODE_LEN_BYTES;
304 let bytes = self.trie_data.get(off..off + NODE_LEN_BYTES)?;
305 Some((
306 LittleEndian::read_u32(&bytes[0..4]),
307 LittleEndian::read_u32(&bytes[4..8]),
308 ))
309 }
310
311 /// Maps a character to its trie code.
312 ///
313 /// # Arguments
314 ///
315 /// * `c` - The character to map.
316 ///
317 /// # Returns
318 ///
319 /// The mapped code, or `None` when the character appears in no key.
320 #[inline(always)]
321 fn map_code(&self, c: char) -> Option<u32> {
322 let ord = c as usize;
323 if ord >= self.table_len {
324 return None;
325 }
326 let off = TABLE_START + ord * 4;
327 let code = LittleEndian::read_u32(self.trie_data.get(off..off + 4)?);
328 (code != INVALID_CODE).then_some(code)
329 }
330
331 /// Looks up the values run for trie key ordinal `key_ord`.
332 ///
333 /// # Arguments
334 ///
335 /// * `key_ord` - The trie value: the key's ordinal in sorted key order.
336 ///
337 /// # Returns
338 ///
339 /// The key's serialized `WordEntry` records, or `None` if either index is
340 /// out of range (malformed input).
341 #[inline(always)]
342 fn entry_bytes(&self, key_ord: u32) -> Option<&[u8]> {
343 let i = key_ord as usize * 4;
344 let idx = self.vals_idx.get(i..i + 8)?;
345 let start = LittleEndian::read_u32(&idx[0..4]) as usize;
346 let end = LittleEndian::read_u32(&idx[4..8]) as usize;
347 self.vals_data
348 .get(start * WordEntry::SERIALIZED_LEN..end * WordEntry::SERIALIZED_LEN)
349 }
350
351 /// Returns the word entries for every key that is a prefix of `chars`.
352 ///
353 /// This is the tokenizer's hot path: one call per lattice-reachable
354 /// position, yielding matches in ascending length order.
355 ///
356 /// # Arguments
357 ///
358 /// * `chars` - The sentence suffix starting at the query position.
359 ///
360 /// # Returns
361 ///
362 /// An iterator of `(serialized WordEntry records, chars consumed)`.
363 #[inline]
364 pub fn common_prefix_search<'a, 'b>(&'a self, chars: &'b [char]) -> CommonPrefixSearch<'a, 'b> {
365 CommonPrefixSearch {
366 dict: self,
367 chars,
368 pos: 0,
369 node_idx: 0,
370 }
371 }
372
373 /// Returns the entries whose surface is a prefix of `s`, with byte-level
374 /// end offsets.
375 ///
376 /// # Arguments
377 ///
378 /// * `s` - The query string.
379 ///
380 /// # Returns
381 ///
382 /// An iterator of `(end byte offset in s, entry)` pairs.
383 pub fn prefix<'a>(&'a self, s: &'a str) -> impl Iterator<Item = (usize, WordEntry)> + 'a {
384 // Match lengths come back in characters; convert through the byte
385 // offset of each char boundary.
386 let boundaries: Vec<usize> = s
387 .char_indices()
388 .map(|(byte_offset, _)| byte_offset)
389 .chain(std::iter::once(s.len()))
390 .collect();
391 let chars: Vec<char> = s.chars().collect();
392 let mut results: Vec<(usize, WordEntry)> = Vec::new();
393 for (entries, end_char) in self.common_prefix_search_owned(&chars) {
394 let end_byte = boundaries[end_char];
395 for chunk in entries.chunks_exact(WordEntry::SERIALIZED_LEN) {
396 results.push((end_byte, WordEntry::deserialize(chunk, true)));
397 }
398 }
399 results.into_iter()
400 }
401
402 /// [`Self::common_prefix_search`] over a temporary char buffer, collecting
403 /// eagerly so the buffer does not need to outlive the iterator.
404 ///
405 /// # Arguments
406 ///
407 /// * `chars` - The query characters.
408 ///
409 /// # Returns
410 ///
411 /// The collected `(entry bytes, chars consumed)` pairs.
412 fn common_prefix_search_owned(&self, chars: &[char]) -> Vec<(&[u8], usize)> {
413 self.common_prefix_search(chars).collect()
414 }
415
416 /// Find `WordEntry`s with surface
417 ///
418 /// # Arguments
419 ///
420 /// * `surface` - The exact surface form to look up.
421 ///
422 /// # Returns
423 ///
424 /// All entries whose surface equals `surface`.
425 pub fn find_surface(&self, surface: &str) -> Vec<WordEntry> {
426 self.find_surface_iter(surface).collect()
427 }
428
429 /// Find `WordEntry`s with surface using lazy evaluation
430 /// This iterator-based approach reduces memory allocations
431 ///
432 /// # Arguments
433 ///
434 /// * `surface` - The exact surface form to look up.
435 ///
436 /// # Returns
437 ///
438 /// An iterator over the matching entries.
439 pub fn find_surface_iter<'a>(
440 &'a self,
441 surface: &'a str,
442 ) -> impl Iterator<Item = WordEntry> + 'a {
443 let chars: Vec<char> = surface.chars().collect();
444 let char_count = chars.len();
445 let mut entries: Vec<WordEntry> = Vec::new();
446 for (bytes, end_char) in self.common_prefix_search_owned(&chars) {
447 if end_char == char_count {
448 for chunk in bytes.chunks_exact(WordEntry::SERIALIZED_LEN) {
449 entries.push(WordEntry::deserialize(chunk, true));
450 }
451 }
452 }
453 entries.into_iter()
454 }
455
456 /// Common prefix iterator using character array input.
457 ///
458 /// `end_char` counts characters consumed from `suffix` -- the natural
459 /// unit for callers that index `&[char]` (the trainer's lattice does
460 /// `pos + end_char`). The retired daachorse implementation returned byte
461 /// lengths here, which silently misplaced edges for any non-ASCII text.
462 ///
463 /// # Arguments
464 ///
465 /// * `suffix` - The sentence suffix to match prefixes of.
466 ///
467 /// # Returns
468 ///
469 /// Matches for every dictionary key that is a prefix of `suffix`.
470 pub fn common_prefix_iterator(&self, suffix: &[char]) -> Vec<Match> {
471 let mut matches = Vec::new();
472 for (bytes, end_char) in self.common_prefix_search(suffix) {
473 for chunk in bytes.chunks_exact(WordEntry::SERIALIZED_LEN) {
474 let word_entry = WordEntry::deserialize(chunk, true);
475 matches.push(Match {
476 word_idx: WordIdx::new(word_entry.word_id().id()),
477 end_char,
478 });
479 }
480 }
481 matches
482 }
483}
484
485/// Iterator over the dictionary keys that are prefixes of a query.
486///
487/// Mirrors crawdad 0.3's `CommonPrefixSearchIter` step for step, but walks
488/// the serialized bytes directly with bounds-checked reads.
489pub struct CommonPrefixSearch<'a, 'b> {
490 /// The dictionary being searched.
491 dict: &'a PrefixDictionary,
492 /// The query characters.
493 chars: &'b [char],
494 /// Characters consumed so far.
495 pos: usize,
496 /// Current trie node.
497 node_idx: u32,
498}
499
500impl<'a> Iterator for CommonPrefixSearch<'a, '_> {
501 type Item = (&'a [u8], usize);
502
503 /// Advances to the next key that is a prefix of the query.
504 ///
505 /// # Returns
506 ///
507 /// The key's serialized `WordEntry` records and the number of characters
508 /// consumed, or `None` when no further prefix matches.
509 #[inline]
510 fn next(&mut self) -> Option<Self::Item> {
511 while self.pos < self.chars.len() {
512 let mc = self.dict.map_code(self.chars[self.pos])?;
513 let (base_raw, _) = self.dict.node(self.node_idx)?;
514 // A leaf stores its value in `base`, not a child offset, so it
515 // has no children and the walk ends.
516 if base_raw & !OFFSET_MASK != 0 {
517 return None;
518 }
519 let child_idx = (base_raw & OFFSET_MASK) ^ mc;
520 let (child_base, child_check) = self.dict.node(child_idx)?;
521 if child_check & OFFSET_MASK != self.node_idx {
522 return None;
523 }
524 self.node_idx = child_idx;
525 self.pos += 1;
526
527 if child_base & !OFFSET_MASK != 0 {
528 // The node itself is a leaf: its base is the value.
529 let entries = self.dict.entry_bytes(child_base & OFFSET_MASK)?;
530 return Some((entries, self.pos));
531 }
532 if child_check & !OFFSET_MASK != 0 {
533 // The node has a leaf child reached by the end marker, whose
534 // code is 0, so the child index is just the base.
535 let leaf_idx = child_base & OFFSET_MASK;
536 let (leaf_base, _) = self.dict.node(leaf_idx)?;
537 let entries = self.dict.entry_bytes(leaf_base & OFFSET_MASK)?;
538 return Some((entries, self.pos));
539 }
540 }
541 None
542 }
543}
544
545/// Builds the "declared size is implausible" load error.
546///
547/// # Returns
548///
549/// A deserialize error describing an arithmetic overflow in the headers.
550fn implausible_size() -> crate::error::LinderaError {
551 LinderaErrorKind::Deserialize
552 .with_error(anyhow::anyhow!("dict.trie declares an implausible size"))
553}
554
555/// rkyv adapter storing a daachorse automaton as its serialized bytes.
556pub struct DoubleArrayArchiver;
557
558impl ArchiveWith<DoubleArrayAhoCorasick<u32>> for DoubleArrayArchiver {
559 type Archived = rkyv::vec::ArchivedVec<u8>;
560 type Resolver = rkyv::vec::VecResolver;
561
562 /// Resolves the archived byte vector for the automaton.
563 fn resolve_with(
564 field: &DoubleArrayAhoCorasick<u32>,
565 resolver: Self::Resolver,
566 out: Place<Self::Archived>,
567 ) {
568 let bytes = field.serialize();
569 rkyv::vec::ArchivedVec::resolve_from_slice(&bytes, resolver, out);
570 }
571}
572
573impl<S: Fallible + rkyv::ser::Writer + rkyv::ser::Allocator + ?Sized>
574 SerializeWith<DoubleArrayAhoCorasick<u32>, S> for DoubleArrayArchiver
575{
576 /// Serializes the automaton as a byte vector.
577 fn serialize_with(
578 field: &DoubleArrayAhoCorasick<u32>,
579 serializer: &mut S,
580 ) -> Result<Self::Resolver, S::Error> {
581 let bytes = field.serialize();
582 rkyv::vec::ArchivedVec::serialize_from_slice(&bytes, serializer)
583 }
584}
585
586impl<D: Fallible<Error: Source> + ?Sized>
587 DeserializeWith<rkyv::vec::ArchivedVec<u8>, DoubleArrayAhoCorasick<u32>, D>
588 for DoubleArrayArchiver
589{
590 /// Deserialize the archived byte vector into a `DoubleArrayAhoCorasick`.
591 ///
592 /// # Returns
593 ///
594 /// The deserialized `DoubleArrayAhoCorasick`, or an error if deserialization fails.
595 fn deserialize_with(
596 archived: &rkyv::vec::ArchivedVec<u8>,
597 _deserializer: &mut D,
598 ) -> Result<DoubleArrayAhoCorasick<u32>, D::Error> {
599 let (da, _) = DoubleArrayAhoCorasick::deserialize(archived.as_slice()).map_err(|err| {
600 D::Error::new(std::io::Error::new(
601 std::io::ErrorKind::InvalidData,
602 err.to_string(),
603 ))
604 })?;
605 Ok(da)
606 }
607}
608
609/// serde adapter storing a daachorse automaton as its serialized bytes.
610mod double_array_serde {
611 use daachorse::DoubleArrayAhoCorasick;
612 use serde::{Deserialize, Deserializer, Serializer};
613
614 /// Serializes the automaton as bytes.
615 ///
616 /// # Arguments
617 ///
618 /// * `da` - The automaton to serialize.
619 /// * `serializer` - The serde serializer.
620 ///
621 /// # Returns
622 ///
623 /// The serializer's output.
624 pub fn serialize<S>(da: &DoubleArrayAhoCorasick<u32>, serializer: S) -> Result<S::Ok, S::Error>
625 where
626 S: Serializer,
627 {
628 let bytes = da.serialize();
629 serializer.serialize_bytes(&bytes)
630 }
631
632 /// Deserializes an automaton from bytes, validating them.
633 ///
634 /// # Arguments
635 ///
636 /// * `deserializer` - The serde deserializer.
637 ///
638 /// # Returns
639 ///
640 /// The automaton, or an error for malformed bytes.
641 pub fn deserialize<'de, D>(deserializer: D) -> Result<DoubleArrayAhoCorasick<u32>, D::Error>
642 where
643 D: Deserializer<'de>,
644 {
645 let bytes: Vec<u8> = Deserialize::deserialize(deserializer)?;
646 let (da, _) = DoubleArrayAhoCorasick::deserialize(&bytes)
647 .map_err(|err| serde::de::Error::custom(err.to_string()))?;
648 Ok(da)
649 }
650}
651
652/// The user prefix dictionary: a daachorse Aho-Corasick automaton.
653///
654/// The field sequence is byte-for-byte the sequence the pre-v6
655/// `PrefixDictionary` archived (`da`, `vals_data`, `words_idx_data`,
656/// `words_data`, `is_system`). rkyv 0.8 archives structurally, without type
657/// names, so keeping the sequence is what lets every previously-built user
658/// dictionary `.bin` keep loading. Do not add, remove or reorder fields
659/// without bumping the dictionary format version and rebuilding the committed
660/// `.bin` fixtures.
661#[derive(Clone, Serialize, Deserialize, Archive, RkyvSerialize, RkyvDeserialize)]
662pub struct UserPrefixDictionary {
663 /// The Aho-Corasick automaton over surface forms.
664 #[serde(with = "self::double_array_serde")]
665 #[rkyv(with = DoubleArrayArchiver)]
666 pub da: DoubleArrayAhoCorasick<u32>,
667 /// The values file: `WordEntry` records back to back.
668 pub vals_data: Data,
669 /// Byte offsets into `words_data`, one `u32` per word id.
670 pub words_idx_data: Data,
671 /// Word detail records.
672 pub words_data: Data,
673 /// Always `false`; retained because the archived field sequence must not
674 /// change (see the type-level comment).
675 pub is_system: bool,
676}
677
678impl UserPrefixDictionary {
679 /// Decode the `(offset, count)` pair stored in the double-array value.
680 ///
681 /// The word-id offset lives in the high 24 bits and the per-surface
682 /// variant count in the low 8 bits (up to 255 variants). The legacy 5-bit
683 /// encoding was retired in v4.0.0; user `.bin` files built with v3 must
684 /// be rebuilt from CSV.
685 ///
686 /// # Arguments
687 ///
688 /// * `val` - The packed automaton value.
689 ///
690 /// # Returns
691 ///
692 /// The `(offset, count)` pair.
693 #[inline]
694 pub fn decode_val(&self, val: u32) -> (u32, u32) {
695 (val >> 8u32, val & ((1u32 << 8) - 1u32))
696 }
697
698 /// Load a `UserPrefixDictionary` from raw binary data.
699 ///
700 /// The automaton bytes are always run through daachorse's validating
701 /// deserializer: user dictionaries come from the filesystem or from
702 /// callers, never from this crate's own build pipeline, so there is no
703 /// trusted path.
704 ///
705 /// # Arguments
706 ///
707 /// * `da_data` - Serialized automaton bytes.
708 /// * `vals_data` - Values data bytes.
709 /// * `words_idx_data` - Word index data bytes.
710 /// * `words_data` - Words data bytes.
711 ///
712 /// # Returns
713 ///
714 /// A `UserPrefixDictionary`, or an error if deserialization fails.
715 pub fn load(
716 da_data: impl Into<Data>,
717 vals_data: impl Into<Data>,
718 words_idx_data: impl Into<Data>,
719 words_data: impl Into<Data>,
720 ) -> LinderaResult<UserPrefixDictionary> {
721 let da_bytes = da_data.into();
722 let da = DoubleArrayAhoCorasick::deserialize(&da_bytes[..])
723 .map_err(|err| {
724 LinderaErrorKind::Deserialize.with_error(anyhow::anyhow!(err.to_string()))
725 })?
726 .0;
727
728 Ok(UserPrefixDictionary {
729 da,
730 vals_data: vals_data.into(),
731 words_idx_data: words_idx_data.into(),
732 words_data: words_data.into(),
733 is_system: false,
734 })
735 }
736}
737
738#[cfg(test)]
739mod tests {
740 use daachorse::DoubleArrayAhoCorasickBuilder;
741
742 use super::*;
743 use crate::viterbi::{LexType, WordId};
744
745 fn entry(word_id: u32, cost: i16) -> WordEntry {
746 WordEntry::new(WordId::new(LexType::System, word_id), cost, 0, 0)
747 }
748
749 fn sample_map() -> BTreeMap<String, Vec<WordEntry>> {
750 let mut map = BTreeMap::new();
751 map.insert("世界".to_string(), vec![entry(0, 10)]);
752 map.insert("世界中".to_string(), vec![entry(1, 20), entry(2, 30)]);
753 map.insert("世論調査".to_string(), vec![entry(3, 40)]);
754 map.insert("統計調査".to_string(), vec![entry(4, 50)]);
755 map
756 }
757
758 /// The load-bearing round-trip: the in-place view must return exactly
759 /// what crawdad's own search returns on the same serialized bytes. This
760 /// is the test that fails loudly if crawdad's byte layout ever changes.
761 #[test]
762 fn trie_view_matches_crawdad_search() {
763 let map = sample_map();
764 let keys: Vec<&str> = map.keys().map(|k| k.as_str()).collect();
765 let reference = crawdad::Trie::from_keys(keys.iter().copied()).unwrap();
766
767 let dict = PrefixDictionary::from_word_entry_map(&map).unwrap();
768
769 for haystack in ["世界中で世論調査", "統計調査だ", "無関係な文", "世", ""]
770 {
771 let chars: Vec<char> = haystack.chars().collect();
772 for start in 0..=chars.len() {
773 let expected: Vec<(u32, usize)> = reference
774 .common_prefix_search(chars[start..].iter().copied())
775 .collect();
776 let actual: Vec<(usize, usize)> = dict
777 .common_prefix_search(&chars[start..])
778 .map(|(bytes, end)| (bytes.len() / WordEntry::SERIALIZED_LEN, end))
779 .collect();
780
781 assert_eq!(actual.len(), expected.len(), "at {haystack:?}[{start}..]");
782 for ((key_ord, exp_end), (n_entries, act_end)) in expected.iter().zip(actual.iter())
783 {
784 assert_eq!(exp_end, act_end);
785 // The run length must match the map's entry count for the
786 // key crawdad says matched.
787 let surface: String = chars[start..start + exp_end].iter().collect();
788 assert_eq!(
789 map[&surface].len(),
790 *n_entries,
791 "run length for key ordinal {key_ord}"
792 );
793 }
794 }
795 }
796 }
797
798 #[test]
799 fn find_surface_returns_all_variants() {
800 let dict = PrefixDictionary::from_word_entry_map(&sample_map()).unwrap();
801
802 let entries = dict.find_surface("世界中");
803 assert_eq!(entries.len(), 2);
804 assert_eq!(entries[0].word_cost(), 20);
805 assert_eq!(entries[1].word_cost(), 30);
806
807 assert!(dict.find_surface("世論").is_empty());
808 assert!(dict.find_surface("未知語").is_empty());
809 }
810
811 #[test]
812 fn prefix_returns_byte_offsets() {
813 let dict = PrefixDictionary::from_word_entry_map(&sample_map()).unwrap();
814
815 let results: Vec<(usize, WordEntry)> = dict.prefix("世界中で").collect();
816 // "世界" (6 bytes) and then "世界中" (9 bytes), three entries total.
817 assert_eq!(results.len(), 3);
818 assert_eq!(results[0].0, 6);
819 assert_eq!(results[1].0, 9);
820 assert_eq!(results[2].0, 9);
821 }
822
823 #[test]
824 fn common_prefix_iterator_counts_characters() {
825 let dict = PrefixDictionary::from_word_entry_map(&sample_map()).unwrap();
826
827 let chars: Vec<char> = "世界中".chars().collect();
828 let matches = dict.common_prefix_iterator(&chars);
829 // "世界" consumes 2 chars, "世界中" consumes 3 -- in characters, not
830 // bytes (the retired implementation returned 6 and 9 here).
831 assert_eq!(matches.len(), 3);
832 assert_eq!(matches[0].end_char, 2);
833 assert_eq!(matches[1].end_char, 3);
834 assert_eq!(matches[2].end_char, 3);
835 }
836
837 #[test]
838 fn load_rejects_truncated_trie() {
839 let map = sample_map();
840 let (trie_bytes, idx_bytes) = PrefixDictionary::serialize_trie(&map).unwrap();
841
842 let mut truncated = trie_bytes.clone();
843 truncated.truncate(trie_bytes.len() - 3);
844 assert!(
845 PrefixDictionary::load(
846 truncated,
847 idx_bytes.clone(),
848 Vec::new(),
849 Vec::new(),
850 Vec::new()
851 )
852 .is_err()
853 );
854
855 assert!(
856 PrefixDictionary::load(vec![0u8; 2], idx_bytes, Vec::new(), Vec::new(), Vec::new())
857 .is_err()
858 );
859 }
860
861 #[test]
862 fn load_rejects_absurd_table_length() {
863 // A table length claiming more entries than the file could hold must
864 // be rejected up front, not fed to an allocator.
865 let mut data = Vec::new();
866 data.extend_from_slice(&u32::MAX.to_le_bytes());
867 data.extend_from_slice(&[0u8; 16]);
868 assert!(
869 PrefixDictionary::load(data, Vec::new(), Vec::new(), Vec::new(), Vec::new()).is_err()
870 );
871 }
872
873 #[test]
874 fn corrupted_nodes_yield_no_matches_without_panicking() {
875 let map = sample_map();
876 let (trie_bytes, idx_bytes) = PrefixDictionary::serialize_trie(&map).unwrap();
877
878 // Flip bytes throughout the node array; every lookup must stay
879 // panic-free (bounds-checked reads make garbage "no match").
880 for step in [1usize, 3, 7, 13] {
881 let mut corrupted = trie_bytes.clone();
882 let start = corrupted.len().saturating_sub(200);
883 let len = corrupted.len();
884 for i in (start..len).step_by(step) {
885 corrupted[i] ^= 0xa5;
886 }
887 if let Ok(dict) = PrefixDictionary::load(
888 corrupted,
889 idx_bytes.clone(),
890 Vec::new(),
891 Vec::new(),
892 Vec::new(),
893 ) {
894 let chars: Vec<char> = "世界中で統計調査".chars().collect();
895 for start in 0..chars.len() {
896 for _ in dict.common_prefix_search(&chars[start..]) {}
897 }
898 }
899 }
900 }
901
902 #[test]
903 fn serialize_trie_rejects_nul_in_surface() {
904 let mut map = BTreeMap::new();
905 map.insert("a\0b".to_string(), vec![entry(0, 0)]);
906 assert!(PrefixDictionary::serialize_trie(&map).is_err());
907 }
908
909 #[test]
910 fn empty_dictionary_matches_nothing() {
911 let dict = PrefixDictionary::from_word_entry_map(&BTreeMap::new()).unwrap();
912 let chars: Vec<char> = "何か".chars().collect();
913 assert_eq!(dict.common_prefix_search(&chars).count(), 0);
914 }
915
916 #[test]
917 fn user_dictionary_load_validates_da_bytes() {
918 let keyset: Vec<(&[u8], u32)> = vec![(b"a", 0), (b"ab", 1), (b"b", 2)];
919 let da = DoubleArrayAhoCorasickBuilder::new()
920 .build_with_values(keyset)
921 .unwrap();
922 let da_bytes = da.serialize();
923
924 let dict = UserPrefixDictionary::load(
925 da_bytes.clone(),
926 Vec::<u8>::new(),
927 Vec::<u8>::new(),
928 Vec::<u8>::new(),
929 )
930 .unwrap();
931 assert_eq!(dict.da.find_overlapping_iter("ab").count(), 3);
932
933 // Truncated bytes must be rejected by the validating deserializer,
934 // not panic or read out of bounds.
935 let mut truncated = da_bytes;
936 truncated.truncate(4);
937 assert!(
938 UserPrefixDictionary::load(
939 truncated,
940 Vec::<u8>::new(),
941 Vec::<u8>::new(),
942 Vec::<u8>::new()
943 )
944 .is_err()
945 );
946 }
947}