sudachi/dic/
lexicon_set.rs

1/*
2 * Copyright (c) 2021-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 thiserror::Error;
18
19use crate::dic::binary_loader::BinaryLexicon;
20use crate::dic::lexicon::strings::StringPointer;
21use crate::dic::lexicon::{Lexicon, LexiconEntry, MAX_DICTIONARIES};
22use crate::dic::subset::InfoSubset;
23use crate::dic::word_id::{DictId, WordId};
24use crate::dic::word_info::{WordInfo, WordInfoEntryIdCursor};
25use crate::dic::LexiconAccess;
26use crate::prelude::*;
27
28/// Sudachi error
29#[derive(Error, Debug, Eq, PartialEq)]
30pub enum LexiconSetError {
31    #[error("too large word_id {0} in dict {1}")]
32    TooLargeWordId(u32, usize),
33
34    #[error("too large dictionary_id {0}")]
35    TooLargeDictionaryId(usize),
36
37    #[error("too many user dictionaries")]
38    TooManyDictionaries,
39
40    #[error("invalid string pointer of length={0}, offset={1}, alignment={2}")]
41    InvalidStringPointer(usize, usize, usize),
42}
43
44/// Set of Lexicons
45///
46/// Handles multiple lexicons as one lexicon
47/// The first lexicon in the list must be from system dictionary
48pub struct LexiconSet<'a> {
49    lexicons: Vec<Lexicon<'a>>,
50    pos_offsets: Vec<usize>,
51    num_system_pos: usize,
52}
53
54#[doc(hidden)]
55pub struct WordIdCursor {
56    lexicon_index: usize,
57    entry_cursor: Option<WordInfoEntryIdCursor>,
58}
59
60impl LexiconAccess for LexiconSet<'_> {
61    fn lexicon(&self) -> &LexiconSet<'_> {
62        self
63    }
64}
65
66impl<'a> LexiconSet<'a> {
67    /// Creates a LexiconSet from a system lexicon
68    pub fn from_system_binary(
69        system_lexicon: BinaryLexicon<'a>,
70        num_system_pos: usize,
71    ) -> LexiconSet<'a> {
72        let mut lexicon = Lexicon::from_binary(system_lexicon);
73        lexicon.set_dic_id(0);
74        LexiconSet {
75            lexicons: vec![lexicon],
76            pos_offsets: vec![0],
77            num_system_pos,
78        }
79    }
80
81    /// Creates a LexiconSet given a system lexicon
82    pub fn new(mut system_lexicon: Lexicon<'a>, num_system_pos: usize) -> LexiconSet<'a> {
83        system_lexicon.set_dic_id(0);
84        LexiconSet {
85            lexicons: vec![system_lexicon],
86            pos_offsets: vec![0],
87            num_system_pos,
88        }
89    }
90
91    /// Add a lexicon to the lexicon list
92    ///
93    /// pos_offset: number of pos in the grammar
94    pub fn append(
95        &mut self,
96        mut lexicon: Lexicon<'a>,
97        pos_offset: usize,
98    ) -> Result<(), LexiconSetError> {
99        if self.is_full() {
100            return Err(LexiconSetError::TooManyDictionaries);
101        }
102        lexicon.set_dic_id(self.lexicons.len() as u8);
103        self.lexicons.push(lexicon);
104        self.pos_offsets.push(pos_offset);
105        Ok(())
106    }
107
108    /// Returns if dictionary capacity is full
109    pub fn is_full(&self) -> bool {
110        self.lexicons.len() >= MAX_DICTIONARIES
111    }
112}
113
114impl LexiconSet<'_> {
115    /// Returns iterator which yields all words in the dictionary, starting from the `offset` bytes
116    ///
117    /// Searches dictionaries in the reverse order: user dictionaries first and then system dictionary
118    #[inline]
119    pub fn lookup<'b>(
120        &'b self,
121        input: &'b [u8],
122        offset: usize,
123    ) -> impl Iterator<Item = LexiconEntry> + 'b {
124        // word_id fixup was moved to lexicon itself
125        self.lexicons
126            .iter()
127            .rev()
128            .flat_map(move |l| l.lookup(input, offset))
129    }
130
131    /// Pipelined + prefetched batch form of [`LexiconSet::lookup`].
132    ///
133    /// Calls `emit(bucket, entry)` for every match; `bucket` indexes `starts`.
134    /// Within a bucket the order matches repeated [`LexiconSet::lookup`] calls
135    /// (user dictionaries first, then system, each in trie-walk order), so
136    /// grouping by `bucket` reproduces the scalar result.
137    #[inline]
138    pub fn lookup_batch<F: FnMut(usize, LexiconEntry)>(
139        &self,
140        input: &[u8],
141        starts: &[usize],
142        mut emit: F,
143    ) {
144        // Reverse dictionary order, one lexicon at a time, to match lookup().
145        for lexicon in self.lexicons.iter().rev() {
146            lexicon.lookup_batch(input, starts, &mut emit);
147        }
148    }
149
150    /// Checks prefix end offsets in the same dictionary order as lookup(), but
151    /// without expanding trie leaves into word IDs.
152    #[inline]
153    pub(crate) fn check_prefix_ends<F>(
154        &self,
155        input: &[u8],
156        offset: usize,
157        mut check: F,
158    ) -> Option<bool>
159    where
160        F: FnMut(usize) -> Option<bool>,
161    {
162        for lexicon in self.lexicons.iter().rev() {
163            for end in lexicon.lookup_prefix_ends(input, offset) {
164                if let Some(result) = check(end) {
165                    return Some(result);
166                }
167            }
168        }
169        None
170    }
171
172    /// Returns WordInfo for given WordId
173    pub fn get_word_info(&self, id: WordId) -> SudachiResult<WordInfo> {
174        self.get_word_info_subset(id, InfoSubset::all())
175    }
176
177    /// Returns WordInfo for given WordId.
178    /// Only fills a requested subset of fields.
179    /// Rest will be of default values (0 or empty).
180    pub fn get_word_info_subset(&self, id: WordId, subset: InfoSubset) -> SudachiResult<WordInfo> {
181        let dict_id = id.dict();
182        let lexicon = self
183            .lexicons
184            .get(dict_id.as_raw() as usize)
185            .ok_or(SudachiError::InvalidWordId(id))?;
186        let word_info_data = lexicon.get_word_info(id.entry(), subset)?.resolve(
187            dict_id,
188            self.num_system_pos,
189            &self.pos_offsets,
190            subset,
191        );
192
193        Ok(WordInfo::new(word_info_data, id))
194    }
195
196    /// Returns word_param for given word_id
197    pub fn get_word_param(&self, id: WordId) -> (i16, i16, i16) {
198        let dict_id = id.dict().as_raw() as usize;
199        self.lexicons[dict_id].get_word_param(id.entry())
200    }
201
202    /// Returns word_param for given word_id.
203    pub fn get_word_param_checked(&self, id: WordId) -> SudachiResult<(i16, i16, i16)> {
204        let dict_id = id.dict().as_raw() as usize;
205        match self.lexicons.get(dict_id) {
206            Some(lexicon) => lexicon
207                .get_word_param_checked(id.entry())
208                .ok_or(SudachiError::InvalidWordId(id)),
209            None => Err(SudachiError::InvalidWordId(id)),
210        }
211    }
212
213    #[inline]
214    pub fn get_string(&self, word_id: WordId, strptr: StringPointer) -> SudachiResult<String> {
215        self.lexicons[word_id.dict().as_raw() as usize].get_string(strptr)
216    }
217
218    pub fn size(&self) -> u32 {
219        self.lexicons.iter().fold(0, |acc, lex| acc + lex.size())
220    }
221
222    pub fn word_ids(&self) -> impl Iterator<Item = SudachiResult<WordId>> + '_ {
223        self.lexicons.iter().enumerate().flat_map(|(dict_id, lex)| {
224            let dict_id = DictId::new(dict_id as u8);
225            lex.entry_ids()
226                .map(move |entry| entry.map(|entry| WordId::from_parts(dict_id, entry)))
227        })
228    }
229
230    #[doc(hidden)]
231    pub fn word_id_cursor(&self) -> WordIdCursor {
232        WordIdCursor {
233            lexicon_index: 0,
234            entry_cursor: self.lexicons.first().map(Lexicon::entry_id_cursor),
235        }
236    }
237
238    #[doc(hidden)]
239    pub fn next_word_id(&self, cursor: &mut WordIdCursor) -> SudachiResult<Option<WordId>> {
240        loop {
241            let Some(lexicon) = self.lexicons.get(cursor.lexicon_index) else {
242                return Ok(None);
243            };
244            let Some(entry_cursor) = cursor.entry_cursor.as_mut() else {
245                return Ok(None);
246            };
247            if let Some(entry) = lexicon.next_entry_id(entry_cursor)? {
248                let dict_id = DictId::new(cursor.lexicon_index as u8);
249                return Ok(Some(WordId::from_parts(dict_id, entry)));
250            }
251
252            cursor.lexicon_index += 1;
253            cursor.entry_cursor = self
254                .lexicons
255                .get(cursor.lexicon_index)
256                .map(Lexicon::entry_id_cursor);
257        }
258    }
259
260    pub fn system_word_ids_in_order(&self) -> Vec<WordId> {
261        if self.lexicons.is_empty() {
262            return Vec::new();
263        }
264        self.lexicons[0]
265            .entry_ids_in_order()
266            .into_iter()
267            .map(|entry| WordId::from_parts(DictId::SYSTEM, entry))
268            .collect()
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use crate::dic::binary_loader::LoadedDictionary;
275
276    const TEST_SYSTEM_DIC: &[u8] = include_bytes!("../../tests/resources/system.dic.test");
277
278    #[test]
279    fn check_prefix_ends_matches_lookup_end_order() {
280        let dictionary = LoadedDictionary::load_system(TEST_SYSTEM_DIC).unwrap();
281        let lexicon_set = &dictionary.lexicon_set;
282        let inputs = [
283            "ばな。なです。",
284            "東京都に行く",
285            "京都",
286            "あいうえお",
287            "1.と2.が。",
288        ];
289
290        for input in inputs {
291            let bytes = input.as_bytes();
292            for (offset, _) in input.char_indices() {
293                let mut checked_ends = Vec::new();
294                let decision = lexicon_set.check_prefix_ends(bytes, offset, |end| {
295                    checked_ends.push(end);
296                    None::<bool>
297                });
298
299                assert_eq!(decision, None);
300
301                let mut expected_ends = Vec::new();
302                for lexicon in lexicon_set.lexicons.iter().rev() {
303                    let mut lookup_ends = Vec::new();
304                    for entry in lexicon.lookup(bytes, offset) {
305                        if lookup_ends.last() != Some(&entry.end) {
306                            lookup_ends.push(entry.end);
307                        }
308                    }
309
310                    let prefix_ends = lexicon
311                        .lookup_prefix_ends(bytes, offset)
312                        .collect::<Vec<_>>();
313                    assert_eq!(lookup_ends, prefix_ends);
314                    expected_ends.extend(prefix_ends);
315                }
316
317                assert_eq!(checked_ends, expected_ends);
318            }
319        }
320    }
321}
sudachi/dic/lexicon_set.rs

sudachi/dic/
lexicon_set.rs
