Skip to main content

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}