Skip to main content

sudachi/dic/
binary_loader.rs

1/*
2 * Copyright (c) 2026 Works Applications Co., Ltd.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17use crate::dic::connect::ConnectionMatrix;
18use crate::dic::description::{Block, Description};
19use crate::dic::error::DictionaryCompatibilityError;
20use crate::dic::grammar::Grammar;
21use crate::dic::header::HeaderError;
22use crate::dic::lexicon::strings::CompactedStrings;
23use crate::dic::lexicon::trie::Trie;
24use crate::dic::lexicon::word_id_table::WordIdTable;
25use crate::dic::lexicon::word_params::WordParams;
26use crate::dic::lexicon::Lexicon;
27use crate::dic::lexicon_set::LexiconSet;
28use crate::dic::pos::PosList;
29use crate::dic::read::utf8_string::utf8_string;
30use crate::dic::read::varint::varint32;
31use crate::dic::word_info::WordInfos;
32use crate::dic::{DescriptionAccess, DictionaryAccess, LexiconAccess, ReferenceIdAccess};
33use crate::plugin::input_text::InputTextPlugin;
34use crate::plugin::oov::OovProviderPlugin;
35use crate::plugin::path_rewrite::PathRewritePlugin;
36use crate::prelude::*;
37use std::collections::HashMap;
38
39/// A single system or user dictionary
40pub struct BinaryDictionary<'a> {
41    raw_bytes: &'a [u8],
42    pub description: Description,
43    pub grammar: BinaryGrammar<'a>,
44    pub lexicon: BinaryLexicon<'a>,
45}
46
47impl<'a> BinaryDictionary<'a> {
48    /// Load a binary dictionary from bytes
49    ///
50    /// # Safety
51    /// This function is marked unsafe because it does not perform header validation
52    unsafe fn load(buf: &'a [u8]) -> SudachiResult<Self> {
53        let description = Description::load(buf)?;
54        let grammar = BinaryGrammar::load(buf, &description)?;
55        let lexicon = BinaryLexicon::load(buf, &description)?;
56
57        Ok(BinaryDictionary {
58            raw_bytes: buf,
59            description,
60            grammar,
61            lexicon,
62        })
63    }
64
65    pub fn load_system(buf: &'a [u8]) -> SudachiResult<Self> {
66        let dict = unsafe { Self::load(buf)? };
67
68        if dict.description.is_system_dictionary() {
69            Ok(dict)
70        } else {
71            // TODO: fix error type
72            Err(SudachiError::InvalidHeader(
73                HeaderError::InvalidSystemDictVersion,
74            ))
75        }
76    }
77
78    pub fn load_user(buf: &'a [u8]) -> SudachiResult<Self> {
79        let dict = unsafe { Self::load(buf)? };
80
81        if dict.description.is_user_dictionary() {
82            Ok(dict)
83        } else {
84            // TODO: fix error type
85            Err(SudachiError::InvalidHeader(
86                HeaderError::InvalidUserDictVersion,
87            ))
88        }
89    }
90
91    pub fn compatibility_key(&self) -> &str {
92        compatibility_key(&self.description)
93    }
94
95    pub fn is_compatible_with(&self, other: &BinaryDictionary<'_>) -> bool {
96        self.compatibility_key() == other.compatibility_key()
97    }
98
99    /// Build-time helper for dictionary builders and dump tooling.
100    /// Runtime tokenization does not use this table.
101    pub fn reference_id_table(&self) -> SudachiResult<HashMap<u32, String>> {
102        let Some(bytes) = self
103            .description
104            .slice_or_none(self.raw_bytes, Block::ReferenceIdTable)?
105        else {
106            return Ok(HashMap::with_capacity(0));
107        };
108        parse_reference_id_table(bytes)
109    }
110}
111
112/// Grammar part of the single binary dictionary
113pub struct BinaryGrammar<'a> {
114    /// The list of part of speechs
115    pub pos_list: PosList,
116
117    /// The overloadable connection cost matrix
118    ///
119    /// Only system dictionary has this.
120    pub connection: Option<ConnectionMatrix<'a>>,
121}
122
123impl<'a> BinaryGrammar<'a> {
124    /// load a grammar from bytes
125    pub fn load(buf: &'a [u8], description: &Description) -> SudachiResult<Self> {
126        let connection_bytes = description.slice_or_none(buf, Block::ConnectionMatrix)?;
127        let connection = match connection_bytes {
128            Some(bytes) => {
129                let connection = ConnectionMatrix::from_bytes(bytes)?;
130                Some(connection)
131            }
132            None => None,
133        };
134
135        let pos_list = PosList::from_bytes(description.slice(buf, Block::POSTable)?)?;
136
137        Ok(Self {
138            pos_list,
139            connection,
140        })
141    }
142}
143
144/// Lexicon part of the single binary dictionary
145pub struct BinaryLexicon<'a> {
146    /// TRIE (double array), mapping from index form to WordIdTable offset
147    pub trie: Trie<'a>,
148    /// list of word ids that have the same index form
149    pub word_id_table: WordIdTable<'a>,
150    /// list of word information (for analysis)
151    pub word_params: WordParams<'a>,
152    /// list of word information (for non-analysis)
153    pub word_infos: WordInfos<'a>,
154    /// Stotage of strings in the lixicon (normalized form etc.)
155    pub strings: CompactedStrings<'a>,
156    /// The number of entries in the lexicon
157    pub num_total_entries: u32,
158}
159
160impl<'a> BinaryLexicon<'a> {
161    /// load a lexicon from bytes
162    pub fn load(buf: &'a [u8], description: &Description) -> SudachiResult<Self> {
163        let trie = Trie::from_bytes(description.slice(buf, Block::TRIEIndex)?);
164        let word_id_table = WordIdTable::from_bytes(description.slice(buf, Block::WordPointers)?);
165
166        // word_params and word_infos share the same byte range.
167        // the first 8 bytes of a word entry is the paramaters and rest is the infos.
168        // handle separately because we use them in different steps; during/after analysis.
169        let entries_bytes = description.slice(buf, Block::Entries)?;
170        let word_params = WordParams::from_bytes(entries_bytes);
171        let word_infos = WordInfos::from_bytes(entries_bytes);
172        word_infos.validate_entry_boundaries(description.num_total_entries())?;
173
174        let strings = CompactedStrings::from_bytes(description.slice(buf, Block::Strings)?);
175
176        Ok(Self {
177            trie,
178            word_id_table,
179            word_params,
180            word_infos,
181            strings,
182            num_total_entries: description.num_total_entries(),
183        })
184    }
185}
186
187/// A dictionary consists of one system_dict and zero or more user_dicts.
188///
189/// This is mostly used for testing purpose.
190pub struct LoadedDictionary<'a> {
191    pub description: Description,
192    pub grammar: Grammar<'a>,
193    pub lexicon_set: LexiconSet<'a>,
194    system_reference_ids: HashMap<u32, String>,
195}
196
197impl<'a> LoadedDictionary<'a> {
198    /// Convert to Loaded dictionary
199    pub fn from_system_binary(binary: BinaryDictionary<'a>) -> SudachiResult<Self> {
200        let system_reference_ids = binary.reference_id_table()?;
201        let description = binary.description;
202        let grammar = Grammar::from_system_binary(binary.grammar)?;
203        let lexicon_set = LexiconSet::from_system_binary(binary.lexicon, grammar.pos_list.len());
204        Ok(LoadedDictionary {
205            description,
206            grammar,
207            lexicon_set,
208            system_reference_ids,
209        })
210    }
211
212    pub fn load_system(bytes: &'a [u8]) -> SudachiResult<Self> {
213        Self::from_system_binary(BinaryDictionary::load_system(bytes)?)
214    }
215
216    pub fn description(&self) -> &Description {
217        &self.description
218    }
219
220    pub fn merge_dictionary(mut self, other: BinaryDictionary<'a>) -> SudachiResult<Self> {
221        let expected_signature = compatibility_key(&self.description);
222        if expected_signature != other.compatibility_key() {
223            return Err(DictionaryCompatibilityError::UserDictionaryWithoutIndex {
224                system_signature: expected_signature.to_owned(),
225                user_reference: other.compatibility_key().to_owned(),
226            }
227            .into());
228        }
229
230        self.lexicon_set.append(
231            Lexicon::from_binary(other.lexicon),
232            self.grammar.pos_list.len(),
233        )?;
234        self.grammar.merge_binary(other.grammar);
235        Ok(self)
236    }
237}
238
239impl LexiconAccess for LoadedDictionary<'_> {
240    fn lexicon(&self) -> &LexiconSet<'_> {
241        &self.lexicon_set
242    }
243}
244
245impl DictionaryAccess for LoadedDictionary<'_> {
246    fn grammar(&self) -> &Grammar<'_> {
247        &self.grammar
248    }
249
250    fn input_text_plugins(&self) -> &[Box<dyn InputTextPlugin + Sync + Send>] {
251        &[]
252    }
253
254    fn oov_provider_plugins(&self) -> &[Box<dyn OovProviderPlugin + Sync + Send>] {
255        &[]
256    }
257
258    fn path_rewrite_plugins(&self) -> &[Box<dyn PathRewritePlugin + Sync + Send>] {
259        &[]
260    }
261}
262
263impl DescriptionAccess for LoadedDictionary<'_> {
264    fn description(&self) -> &Description {
265        &self.description
266    }
267}
268
269impl ReferenceIdAccess for BinaryDictionary<'_> {
270    fn reference_ids(&self) -> HashMap<u32, String> {
271        self.reference_id_table().unwrap_or_default()
272    }
273}
274
275impl ReferenceIdAccess for LoadedDictionary<'_> {
276    fn reference_ids(&self) -> HashMap<u32, String> {
277        self.system_reference_ids.clone()
278    }
279}
280
281fn compatibility_key(description: &Description) -> &str {
282    if description.is_system_dictionary() {
283        description.signature()
284    } else {
285        description.reference()
286    }
287}
288
289fn parse_reference_id_table(bytes: &[u8]) -> SudachiResult<HashMap<u32, String>> {
290    let (mut rest, count) = varint32(bytes)?;
291    let mut result = HashMap::with_capacity(count as usize);
292    for _ in 0..count {
293        let (new_rest, entry_id) = varint32(rest)?;
294        let (new_rest, reference_id) = utf8_string(new_rest)?;
295        result.insert(entry_id, reference_id);
296        rest = new_rest;
297    }
298    Ok(result)
299}