Skip to main content

sudachi/dic/
grammar.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 itertools::Itertools;
18use std::ops::Index;
19
20use crate::dic::binary_loader::BinaryGrammar;
21use crate::dic::character_category::CharacterCategory;
22use crate::dic::connect::ConnectionMatrix;
23use crate::dic::pos::{PosList, POS_DEPTH};
24use crate::prelude::*;
25
26/// Dictionary grammar
27///
28/// Contains part_of_speech list and connection cost map.
29/// It also holds character category.
30pub struct Grammar<'a> {
31    /// The list of part of speechs used in the dictionary
32    pub pos_list: PosList,
33
34    /// The mapping to overload cost table
35    connection: ConnectionMatrix<'a>,
36
37    /// The mapping from character to character_category_type
38    pub character_category: CharacterCategory,
39}
40
41impl<'a> Grammar<'a> {
42    pub const INHIBITED_CONNECTION: i16 = i16::MAX;
43
44    pub const BOS_PARAMETER: (i16, i16, i16) = (0, 0, 0); // left_id, right_id, cost
45    pub const EOS_PARAMETER: (i16, i16, i16) = (0, 0, 0); // left_id, right_id, cost
46
47    pub fn from_system_binary(binary_grammar: BinaryGrammar<'a>) -> SudachiResult<Grammar<'a>> {
48        let connection = binary_grammar
49            .connection
50            .ok_or(SudachiError::ConnectionMatrixMissing)?;
51
52        Ok(Self::from_parts(binary_grammar.pos_list, connection))
53    }
54
55    pub(crate) fn from_parts(pos_list: PosList, connection: ConnectionMatrix<'a>) -> Self {
56        Grammar {
57            pos_list,
58            connection,
59            character_category: CharacterCategory::default(),
60        }
61    }
62
63    /// Merge a another (user) grammar into this grammar
64    ///
65    /// Only pos_list is merged
66    pub fn merge(&mut self, other: Grammar) {
67        self.pos_list.extend(other.pos_list);
68    }
69
70    /// Merge a another (user) binary grammar into this grammar
71    ///
72    /// Only pos_list is merged
73    pub fn merge_binary(&mut self, other: BinaryGrammar) {
74        self.pos_list.extend(other.pos_list);
75    }
76
77    /// Returns connection cost of nodes
78    ///
79    /// left_id: right_id of left node
80    /// right_id: left_if of right node
81    #[inline(always)]
82    pub fn connect_cost(&self, left_id: i16, right_id: i16) -> i16 {
83        self.connection.cost(left_id as u16, right_id as u16)
84    }
85
86    #[inline]
87    pub fn conn_matrix(&self) -> &ConnectionMatrix<'_> {
88        &self.connection
89    }
90
91    /// Sets character category
92    ///
93    /// This is the only way to set character category.
94    /// Character category will be a empty map by default.
95    pub fn set_character_category(&mut self, character_category: CharacterCategory) {
96        self.character_category = character_category;
97    }
98
99    /// Sets connect cost for a specific pair of ids
100    ///
101    /// left_id: right_id of left node
102    /// right_id: left_if of right node
103    pub fn set_connect_cost(&mut self, left_id: i16, right_id: i16, cost: i16) {
104        // for edit connection cost plugin
105        self.connection
106            .update(left_id as u16, right_id as u16, cost);
107    }
108
109    /// Returns a pos_id of given pos in the grammar
110    pub fn get_part_of_speech_id<S>(&self, pos1: &[S]) -> Option<u16>
111    where
112        S: AsRef<str>,
113    {
114        if pos1.len() != POS_DEPTH {
115            return None;
116        }
117        for (i, pos2) in self.pos_list.iter().enumerate() {
118            if pos1.iter().zip(pos2).all(|(a, b)| a.as_ref() == b) {
119                return Some(i as u16);
120            }
121        }
122        None
123    }
124
125    pub fn register_pos<S>(&mut self, pos: &[S]) -> SudachiResult<u16>
126    where
127        S: AsRef<str> + ToString,
128    {
129        if pos.len() != POS_DEPTH {
130            let pos_string = pos.iter().map(|x| x.as_ref()).join(",");
131            return Err(SudachiError::InvalidPartOfSpeech(pos_string));
132        }
133        match self.get_part_of_speech_id(pos) {
134            Some(id) => Ok(id),
135            None => {
136                let new_id = self.pos_list.len();
137                if new_id > u16::MAX as usize {
138                    return Err(SudachiError::InvalidPartOfSpeech(
139                        "Too much POS tags registered".to_owned(),
140                    ));
141                }
142                let components = pos.iter().map(|x| x.to_string()).collect();
143                self.pos_list.push(components);
144                Ok(new_id as u16)
145            }
146        }
147    }
148
149    /// Gets POS components for POS ID.
150    /// Panics if out of bounds.
151    pub fn pos_components(&self, pos_id: u16) -> &[String] {
152        self.pos_list.index(pos_id as usize)
153    }
154}
155
156impl Grammar<'static> {
157    pub fn empty() -> Self {
158        Self::from_parts(PosList::default(), ConnectionMatrix::empty())
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165
166    #[test]
167    fn storage_size() {
168        let grammar = setup_grammar();
169        assert_eq!(grammar.pos_list.len(), 3);
170        assert_eq!(grammar.conn_matrix().num_left(), 3);
171        assert_eq!(grammar.conn_matrix().num_right(), 3);
172    }
173
174    #[test]
175    fn partofspeech_string() {
176        let grammar = setup_grammar();
177        assert_eq!(6, grammar.pos_list[0].len());
178        assert_eq!("BOS/EOS", grammar.pos_list[0][0]);
179        assert_eq!("*", grammar.pos_list[0][5]);
180
181        assert_eq!("一般", grammar.pos_list[1][1]);
182        assert_eq!("*", grammar.pos_list[1][5]);
183
184        assert_eq!("五段-サ行", grammar.pos_list[2][4]);
185        assert_eq!("終止形-一般", grammar.pos_list[2][5]);
186    }
187
188    #[test]
189    fn get_connect_cost() {
190        let grammar = setup_grammar();
191        assert_eq!(0, grammar.connect_cost(0, 0));
192        assert_eq!(-100, grammar.connect_cost(2, 1));
193        assert_eq!(200, grammar.connect_cost(1, 2));
194    }
195
196    #[test]
197    fn set_connect_cost() {
198        let mut grammar = setup_grammar();
199        grammar.set_connect_cost(0, 0, 300);
200        assert_eq!(300, grammar.connect_cost(0, 0));
201    }
202
203    #[test]
204    fn register_pos() {
205        let mut grammar = setup_grammar();
206
207        let id1 = grammar
208            .register_pos(["a", "b", "c", "d", "e", "f"].as_slice())
209            .expect("failed");
210        let id2 = grammar
211            .register_pos(["a", "b", "c", "d", "e", "f"].as_slice())
212            .expect("failed");
213        assert_eq!(id1, id2);
214    }
215
216    #[test]
217    fn bos_parameter() {
218        assert_eq!(0, Grammar::BOS_PARAMETER.0);
219        assert_eq!(0, Grammar::BOS_PARAMETER.1);
220        assert_eq!(0, Grammar::BOS_PARAMETER.2);
221    }
222
223    #[test]
224    fn eos_parameter() {
225        assert_eq!(0, Grammar::EOS_PARAMETER.0);
226        assert_eq!(0, Grammar::EOS_PARAMETER.1);
227        assert_eq!(0, Grammar::EOS_PARAMETER.2);
228    }
229
230    fn setup_grammar() -> Grammar<'static> {
231        let mut pos_list = PosList::default();
232        let mut conn_bytes: Vec<u8> = Vec::new();
233        build_connect_table(&mut conn_bytes);
234        build_part_of_speech(&mut pos_list);
235        let connection = ConnectionMatrix::from_bytes(Box::leak(conn_bytes.into_boxed_slice()))
236            .expect("failed to create conn");
237        Grammar {
238            pos_list,
239            connection,
240            character_category: CharacterCategory::default(),
241        }
242    }
243
244    fn build_part_of_speech(pos_list: &mut PosList) {
245        pos_list.push(vec![
246            "BOS/EOS".to_string(),
247            "*".to_string(),
248            "*".to_string(),
249            "*".to_string(),
250            "*".to_string(),
251            "*".to_string(),
252        ]);
253        pos_list.push(vec![
254            "名詞".to_string(),
255            "一般".to_string(),
256            "*".to_string(),
257            "*".to_string(),
258            "*".to_string(),
259            "*".to_string(),
260        ]);
261        pos_list.push(vec![
262            "動詞".to_string(),
263            "一般".to_string(),
264            "*".to_string(),
265            "*".to_string(),
266            "五段-サ行".to_string(),
267            "終止形-一般".to_string(),
268        ]);
269    }
270    fn build_connect_table(storage: &mut Vec<u8>) {
271        storage.extend(&3_i16.to_le_bytes());
272        storage.extend(&3_i16.to_le_bytes());
273
274        storage.extend(&0_i16.to_le_bytes());
275        storage.extend(&(-300_i16).to_le_bytes());
276        storage.extend(&300_i16.to_le_bytes());
277
278        storage.extend(&300_i16.to_le_bytes());
279        storage.extend(&(-500_i16).to_le_bytes());
280        storage.extend(&(-100_i16).to_le_bytes());
281
282        storage.extend(&(-3000_i16).to_le_bytes());
283        storage.extend(&200_i16.to_le_bytes());
284        storage.extend(&2000_i16.to_le_bytes());
285    }
286}