sudachi/dic/
lexicon.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 std::cmp;
18
19use self::trie::Trie;
20use self::word_id_table::WordIdTable;
21use self::word_params::WordParams;
22use crate::analysis::stateful_tokenizer::StatefulTokenizer;
23use crate::dic::binary_loader::BinaryLexicon;
24use crate::dic::lexicon::strings::CompactedStrings;
25use crate::dic::subset::InfoSubset;
26use crate::dic::word_id::{EntryId, WordId};
27use crate::dic::word_info::{WordInfoEntryIdCursor, WordInfoRefData, WordInfos};
28use crate::dic::DictionaryAccess;
29use crate::prelude::*;
30
31pub mod strings;
32pub mod trie;
33pub mod word_id_table;
34pub mod word_infos;
35pub mod word_params;
36
37/// The first 4 bits of word_id are used to indicate that from which lexicon
38/// the word comes, thus we can only hold 15 lexicons in the same time.
39/// 16th is reserved for marking OOVs.
40pub const MAX_DICTIONARIES: usize = 15;
41
42/// Dictionary lexicon
43///
44/// Contains trie, word_id, word_param, word_info
45pub struct Lexicon<'a> {
46    lex_id: u8,
47
48    trie: Trie<'a>,
49    word_id_table: WordIdTable<'a>,
50    word_params: WordParams<'a>,
51    word_infos: WordInfos<'a>,
52    strings: CompactedStrings<'a>,
53
54    num_total_entries: u32,
55}
56
57impl<'a> Lexicon<'a> {
58    const USER_DICT_COST_PER_MORPH: i32 = -20;
59
60    pub fn from_binary(binary_lexicon: BinaryLexicon<'a>) -> Self {
61        Self {
62            trie: binary_lexicon.trie,
63            word_id_table: binary_lexicon.word_id_table,
64            word_params: binary_lexicon.word_params,
65            word_infos: binary_lexicon.word_infos,
66            strings: binary_lexicon.strings,
67            lex_id: u8::MAX,
68            num_total_entries: binary_lexicon.num_total_entries,
69        }
70    }
71
72    /// Returns the number of entries in the lexicon
73    pub fn size(&self) -> u32 {
74        self.num_total_entries
75    }
76
77    pub fn entry_ids_in_order(&self) -> Vec<EntryId> {
78        match self.word_infos.entry_ids_in_order(self.num_total_entries) {
79            Some(result) => result,
80            None => {
81                // Fallback to trie-indexed entries for malformed binaries.
82                let mut result: Vec<EntryId> = self.word_id_table.all_entries().collect();
83                result.sort_unstable();
84                result
85            }
86        }
87    }
88
89    pub(crate) fn entry_ids(&self) -> impl Iterator<Item = SudachiResult<EntryId>> + '_ {
90        self.word_infos.entry_ids(self.num_total_entries)
91    }
92
93    pub(crate) fn entry_id_cursor(&self) -> WordInfoEntryIdCursor {
94        WordInfos::entry_id_cursor(self.num_total_entries)
95    }
96
97    pub(crate) fn next_entry_id(
98        &self,
99        cursor: &mut WordInfoEntryIdCursor,
100    ) -> SudachiResult<Option<EntryId>> {
101        self.word_infos.next_entry_id(cursor)
102    }
103
104    /// Assign lexicon id to the current Lexicon
105    pub fn set_dic_id(&mut self, id: u8) {
106        assert!(id < MAX_DICTIONARIES as u8);
107        self.lex_id = id
108    }
109
110    #[inline]
111    fn word_id(&self, entry_id: u32) -> WordId {
112        WordId::new(self.lex_id, entry_id)
113    }
114
115    /// Returns an iterator of word_id and end of words that matches given input
116    #[inline]
117    pub fn lookup(
118        &'a self,
119        input: &'a [u8],
120        offset: usize,
121    ) -> impl Iterator<Item = LexiconEntry> + 'a {
122        debug_assert!(self.lex_id < MAX_DICTIONARIES as u8);
123        self.trie
124            .common_prefix_iterator(input, offset)
125            .flat_map(move |e| {
126                self.word_id_table
127                    .entries(e.value as usize)
128                    .map(move |eid| LexiconEntry::new(self.word_id(eid.as_raw()), e.end))
129            })
130    }
131
132    /// Pipelined + prefetched batch form of [`Lexicon::lookup`].
133    ///
134    /// Calls `emit(bucket, entry)` for every match; `bucket` indexes `starts`.
135    /// Within a bucket the order matches [`Lexicon::lookup`], so grouping by
136    /// `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        debug_assert!(self.lex_id < MAX_DICTIONARIES as u8);
145        self.trie
146            .common_prefix_batch(input, starts, |bucket, value, end| {
147                for eid in self.word_id_table.entries(value as usize) {
148                    emit(bucket, LexiconEntry::new(self.word_id(eid.as_raw()), end));
149                }
150            });
151    }
152
153    /// Returns end offsets of trie prefixes that match given input.
154    #[inline]
155    pub(crate) fn lookup_prefix_ends(
156        &'a self,
157        input: &'a [u8],
158        offset: usize,
159    ) -> impl Iterator<Item = usize> + 'a {
160        self.trie
161            .common_prefix_iterator(input, offset)
162            .map(|entry| entry.end)
163    }
164
165    /// Returns WordInfo for given word_id
166    ///
167    /// WordInfo will contain only fields included in InfoSubset
168    pub fn get_word_info(
169        &self,
170        entry_id: EntryId,
171        subset: InfoSubset,
172    ) -> SudachiResult<WordInfoRefData> {
173        self.word_infos.get_word_info(entry_id, subset)
174    }
175
176    /// Returns word_param for given word_id.
177    /// Params are (left_id, right_id, cost).
178    #[inline]
179    pub fn get_word_param(&self, entry_id: EntryId) -> (i16, i16, i16) {
180        let params = self.word_params.get_params(entry_id);
181        (params.left_id(), params.right_id(), params.cost())
182    }
183
184    pub fn get_word_param_checked(&self, entry_id: EntryId) -> Option<(i16, i16, i16)> {
185        let params = self.word_params.get_params_checked(entry_id)?;
186        Some((params.left_id(), params.right_id(), params.cost()))
187    }
188
189    #[inline]
190    pub fn get_string(&self, strptr: strings::StringPointer) -> SudachiResult<String> {
191        self.strings.get_string(strptr)
192    }
193
194    /// update word_param cost based on current tokenizer
195    pub fn update_cost<D: DictionaryAccess>(&mut self, dict: &D) -> SudachiResult<()> {
196        let mut tok = StatefulTokenizer::create(dict, false, Mode::C);
197        let mut ms = MorphemeList::empty(dict);
198
199        for entry_id in self.word_id_table.all_entries() {
200            if self.word_params.get_cost(entry_id) != i16::MIN {
201                continue;
202            }
203            // headword does not requires resolution
204            let wi = self.get_word_info(entry_id, InfoSubset::HEADWORD)?;
205            tok.reset()
206                .push_str(self.strings.get_string(wi.headword_strptr())?.as_str());
207            tok.do_tokenize()?;
208            ms.collect_results(&mut tok)?;
209            let internal_cost = ms.get_internal_cost();
210            let cost = internal_cost + Lexicon::USER_DICT_COST_PER_MORPH * ms.len() as i32;
211            let cost = cmp::min(cost, i16::MAX as i32);
212            let cost = cmp::max(cost, i16::MIN as i32);
213            self.word_params.set_cost(entry_id, cost as i16);
214        }
215
216        Ok(())
217    }
218}
219
220/// Result of the Lexicon lookup
221#[derive(Eq, PartialEq, Debug)]
222pub struct LexiconEntry {
223    /// Id of the returned word
224    pub word_id: WordId,
225    /// Byte index of the word end
226    pub end: usize,
227}
228
229impl LexiconEntry {
230    pub fn new(word_id: WordId, end: usize) -> LexiconEntry {
231        LexiconEntry { word_id, end }
232    }
233}
sudachi/dic/lexicon.rs

sudachi/dic/
lexicon.rs
