1use 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
26pub struct Grammar<'a> {
31 pub pos_list: PosList,
33
34 connection: ConnectionMatrix<'a>,
36
37 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); pub const EOS_PARAMETER: (i16, i16, i16) = (0, 0, 0); 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 pub fn merge(&mut self, other: Grammar) {
67 self.pos_list.extend(other.pos_list);
68 }
69
70 pub fn merge_binary(&mut self, other: BinaryGrammar) {
74 self.pos_list.extend(other.pos_list);
75 }
76
77 #[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 pub fn set_character_category(&mut self, character_category: CharacterCategory) {
96 self.character_category = character_category;
97 }
98
99 pub fn set_connect_cost(&mut self, left_id: i16, right_id: i16, cost: i16) {
104 self.connection
106 .update(left_id as u16, right_id as u16, cost);
107 }
108
109 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 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}