lindera_dictionary/builder/
connection_cost_matrix.rs1use std::borrow::Cow;
2use std::fs::File;
3use std::io::{self, BufRead, Write};
4use std::path::Path;
5use std::sync::Arc;
6
7use encoding_rs::{Encoding, UTF_16BE, UTF_16LE};
8use log::debug;
9use memchr::memchr;
10
11use crate::LinderaResult;
12use crate::dictionary::context_id_map::ContextIdMap;
13use crate::error::LinderaErrorKind;
14use crate::util::{read_file, write_data};
15
16const UTF8_BOM: &[u8] = &[0xEF, 0xBB, 0xBF];
20
21#[cfg(not(target_family = "wasm"))]
26const PARALLEL_THRESHOLD: usize = 1 << 20; #[derive(Debug)]
30pub struct ConnectionCostMatrixBuilder {
31 encoding: Cow<'static, str>,
35 context_id_remap: Option<Arc<ContextIdMap>>,
41}
42
43#[derive(Debug, Default)]
46pub struct ConnectionCostMatrixBuilderOptions {
47 encoding: Option<Cow<'static, str>>,
48 context_id_remap: Option<Arc<ContextIdMap>>,
49}
50
51impl ConnectionCostMatrixBuilderOptions {
52 pub fn encoding(&mut self, value: impl Into<Cow<'static, str>>) -> &mut Self {
53 self.encoding = Some(value.into());
54 self
55 }
56
57 pub fn context_id_remap(&mut self, value: Option<Arc<ContextIdMap>>) -> &mut Self {
58 self.context_id_remap = value;
59 self
60 }
61
62 pub fn builder(&self) -> ConnectionCostMatrixBuilder {
63 ConnectionCostMatrixBuilder {
64 encoding: self.encoding.clone().unwrap_or_else(|| "UTF-8".into()),
65 context_id_remap: self.context_id_remap.clone(),
66 }
67 }
68}
69
70impl ConnectionCostMatrixBuilder {
71 pub fn build(&self, input_dir: &Path, output_dir: &Path) -> LinderaResult<()> {
89 let matrix_data_path = input_dir.join("matrix.def");
90 debug!("reading {matrix_data_path:?}");
91 let buffer = read_file(&matrix_data_path)?;
92
93 let decoded = self.decode_if_needed(&buffer)?;
96 let bytes: &[u8] = match &decoded {
97 Some(decoded) => decoded.as_bytes(),
98 None => strip_utf8_bom(&buffer),
99 };
100
101 let header_end = memchr(b'\n', bytes).unwrap_or(bytes.len());
103 let mut header_pos = 0;
104 let forward_size = next_int(&bytes[..header_end], &mut header_pos).ok_or_else(|| {
105 LinderaErrorKind::Content
106 .with_error(anyhow::anyhow!("matrix.def is missing the size header"))
107 })? as u32;
108 let backward_size = next_int(&bytes[..header_end], &mut header_pos).ok_or_else(|| {
109 LinderaErrorKind::Content.with_error(anyhow::anyhow!(
110 "matrix.def header is missing backward size"
111 ))
112 })? as u32;
113
114 if let Some(remap) = self.context_id_remap.as_deref()
118 && (remap.right.len() != forward_size as usize
119 || remap.left.len() != backward_size as usize)
120 {
121 return Err(LinderaErrorKind::Content.with_error(anyhow::anyhow!(
122 "context-id remap size mismatch: remap.right={} vs forward_size={}, remap.left={} vs backward_size={}",
123 remap.right.len(),
124 forward_size,
125 remap.left.len(),
126 backward_size
127 )));
128 }
129
130 let len = 3 + (forward_size as usize) * (backward_size as usize);
131 let mut costs = vec![i16::MAX; len];
132 costs[0] = -1; costs[1] = forward_size as i16;
134 costs[2] = backward_size as i16;
135
136 let data = if header_end < bytes.len() {
140 &bytes[header_end + 1..]
141 } else {
142 &[]
143 };
144 self.fill_costs(data, forward_size, &mut costs)?;
145
146 let mut matrix_mtx_buffer = Vec::with_capacity(costs.len() * 2);
149 for cost in &costs {
150 matrix_mtx_buffer.extend_from_slice(&cost.to_le_bytes());
151 }
152
153 let wtr_matrix_mtx_path = output_dir.join(Path::new("matrix.mtx"));
154 let mut wtr_matrix_mtx = io::BufWriter::new(
155 File::create(wtr_matrix_mtx_path)
156 .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?,
157 );
158 write_data(&matrix_mtx_buffer, &mut wtr_matrix_mtx)?;
159 wtr_matrix_mtx
160 .flush()
161 .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?;
162
163 Ok(())
164 }
165
166 fn decode_if_needed(&self, buffer: &[u8]) -> LinderaResult<Option<String>> {
179 let encoding =
180 Encoding::for_label_no_replacement(self.encoding.as_bytes()).ok_or_else(|| {
181 LinderaErrorKind::Decode
182 .with_error(anyhow::anyhow!("Invalid encoding: {}", self.encoding))
183 })?;
184
185 let is_utf16 = encoding == UTF_16LE || encoding == UTF_16BE || has_utf16_bom(buffer);
186 if is_utf16 {
187 Ok(Some(encoding.decode(buffer).0.into_owned()))
190 } else {
191 Ok(None)
192 }
193 }
194
195 #[cfg(not(target_family = "wasm"))]
205 fn fill_costs(&self, data: &[u8], forward_size: u32, costs: &mut [i16]) -> LinderaResult<()> {
206 let remap = self.context_id_remap.as_deref();
207 if data.len() >= PARALLEL_THRESHOLD {
208 fill_costs_parallel(data, forward_size, costs, remap)
209 } else {
210 fill_costs_sequential(data, forward_size, costs, remap)
211 }
212 }
213
214 #[cfg(target_family = "wasm")]
223 fn fill_costs(&self, data: &[u8], forward_size: u32, costs: &mut [i16]) -> LinderaResult<()> {
224 fill_costs_sequential(data, forward_size, costs, self.context_id_remap.as_deref())
225 }
226}
227
228fn strip_utf8_bom(buffer: &[u8]) -> &[u8] {
238 buffer.strip_prefix(UTF8_BOM).unwrap_or(buffer)
239}
240
241fn has_utf16_bom(buffer: &[u8]) -> bool {
253 buffer.starts_with(&[0xFF, 0xFE]) || buffer.starts_with(&[0xFE, 0xFF])
254}
255
256pub(crate) fn read_matrix_header(input_dir: &Path, _encoding: &str) -> LinderaResult<(u32, u32)> {
275 let path = input_dir.join("matrix.def");
276 let file = File::open(&path).map_err(|err| {
277 LinderaErrorKind::Io
278 .with_error(anyhow::anyhow!(err))
279 .add_context(format!("Failed to open matrix.def: {path:?}"))
280 })?;
281 let mut reader = io::BufReader::new(file);
282 let mut line = Vec::new();
283 reader.read_until(b'\n', &mut line).map_err(|err| {
284 LinderaErrorKind::Io
285 .with_error(anyhow::anyhow!(err))
286 .add_context("Failed to read matrix.def header line")
287 })?;
288 let bytes = strip_utf8_bom(&line);
289 let mut pos = 0;
290 let forward_size = next_int(bytes, &mut pos).ok_or_else(|| {
291 LinderaErrorKind::Content
292 .with_error(anyhow::anyhow!("matrix.def is missing the size header"))
293 })? as u32;
294 let backward_size = next_int(bytes, &mut pos).ok_or_else(|| {
295 LinderaErrorKind::Content.with_error(anyhow::anyhow!(
296 "matrix.def header is missing backward size"
297 ))
298 })? as u32;
299 Ok((forward_size, backward_size))
300}
301
302fn next_int(bytes: &[u8], pos: &mut usize) -> Option<i32> {
315 while *pos < bytes.len() && matches!(bytes[*pos], b' ' | b'\t' | b'\r') {
316 *pos += 1;
317 }
318 if *pos >= bytes.len() {
319 return None;
320 }
321 let negative = bytes[*pos] == b'-';
322 if negative {
323 *pos += 1;
324 }
325 let start = *pos;
326 let mut value: i32 = 0;
327 while *pos < bytes.len() && bytes[*pos].is_ascii_digit() {
328 value = value
332 .wrapping_mul(10)
333 .wrapping_add((bytes[*pos] - b'0') as i32);
334 *pos += 1;
335 }
336 if *pos == start {
337 return None;
339 }
340 Some(if negative { -value } else { value })
341}
342
343fn parse_data_line(
360 line: &[u8],
361 forward_size: u32,
362 costs_len: usize,
363 remap: Option<&ContextIdMap>,
364) -> LinderaResult<Option<(usize, i16)>> {
365 let mut pos = 0;
366 let Some(forward_id) = next_int(line, &mut pos) else {
367 return Ok(None);
369 };
370 let backward_id = next_int(line, &mut pos).ok_or_else(|| {
371 LinderaErrorKind::Content
372 .with_error(anyhow::anyhow!("matrix.def line is missing backward id"))
373 })?;
374 let cost = next_int(line, &mut pos).ok_or_else(|| {
375 LinderaErrorKind::Content.with_error(anyhow::anyhow!("matrix.def line is missing cost"))
376 })?;
377
378 let forward_id = forward_id as u32 as usize;
379 let backward_id = backward_id as u32 as usize;
380 let (fwd, bwd) = match remap {
385 Some(m) => {
386 if forward_id >= m.right.len() || backward_id >= m.left.len() {
387 return Err(LinderaErrorKind::Content.with_error(anyhow::anyhow!(
388 "matrix.def entry ({forward_id}, {backward_id}) is out of range"
389 )));
390 }
391 (m.right[forward_id] as usize, m.left[backward_id] as usize)
392 }
393 None => (forward_id, backward_id),
394 };
395 let index = 3 + fwd + bwd * forward_size as usize;
396 if index >= costs_len {
397 return Err(LinderaErrorKind::Content.with_error(anyhow::anyhow!(
398 "matrix.def entry ({forward_id}, {backward_id}) is out of range"
399 )));
400 }
401 let cost = (cost as u16) as i16;
402 Ok(Some((index, cost)))
403}
404
405fn fill_costs_sequential(
413 data: &[u8],
414 forward_size: u32,
415 costs: &mut [i16],
416 remap: Option<&ContextIdMap>,
417) -> LinderaResult<()> {
418 let costs_len = costs.len();
419 let mut pos = 0;
420 while pos < data.len() {
421 let line_end = memchr(b'\n', &data[pos..])
422 .map(|offset| pos + offset)
423 .unwrap_or(data.len());
424 if let Some((index, cost)) =
425 parse_data_line(&data[pos..line_end], forward_size, costs_len, remap)?
426 {
427 costs[index] = cost;
428 }
429 pos = line_end + 1;
430 }
431 Ok(())
432}
433
434#[cfg(not(target_family = "wasm"))]
446fn fill_costs_parallel(
447 data: &[u8],
448 forward_size: u32,
449 costs: &mut [i16],
450 remap: Option<&ContextIdMap>,
451) -> LinderaResult<()> {
452 use rayon::prelude::*;
453
454 let costs_len = costs.len();
455 let n_chunks = (rayon::current_num_threads() * 4).max(1);
456
457 let mut bounds = Vec::with_capacity(n_chunks + 1);
459 bounds.push(0usize);
460 for i in 1..n_chunks {
461 let target = data.len() * i / n_chunks;
462 let last = *bounds.last().unwrap_or(&0);
463 if target <= last {
464 continue;
465 }
466 if let Some(offset) = memchr(b'\n', &data[target..]) {
467 let boundary = target + offset + 1;
468 if boundary > last && boundary < data.len() {
469 bounds.push(boundary);
470 }
471 }
472 }
473 bounds.push(data.len());
474
475 let chunks: Vec<&[u8]> = bounds.windows(2).map(|w| &data[w[0]..w[1]]).collect();
476 let partials: Vec<Vec<(usize, i16)>> = chunks
477 .par_iter()
478 .map(|chunk| parse_chunk(chunk, forward_size, costs_len, remap))
479 .collect::<LinderaResult<Vec<_>>>()?;
480
481 for partial in &partials {
482 for &(index, cost) in partial {
483 costs[index] = cost;
484 }
485 }
486 Ok(())
487}
488
489#[cfg(not(target_family = "wasm"))]
501fn parse_chunk(
502 chunk: &[u8],
503 forward_size: u32,
504 costs_len: usize,
505 remap: Option<&ContextIdMap>,
506) -> LinderaResult<Vec<(usize, i16)>> {
507 let mut out = Vec::with_capacity(chunk.len() / 8);
508 let mut pos = 0;
509 while pos < chunk.len() {
510 let line_end = memchr(b'\n', &chunk[pos..])
511 .map(|offset| pos + offset)
512 .unwrap_or(chunk.len());
513 if let Some(entry) = parse_data_line(&chunk[pos..line_end], forward_size, costs_len, remap)?
514 {
515 out.push(entry);
516 }
517 pos = line_end + 1;
518 }
519 Ok(out)
520}
521
522#[cfg(test)]
523mod tests {
524 use super::*;
525
526 fn reference_costs(matrix: &str) -> Vec<i16> {
529 let mut lines = Vec::new();
530 for line in matrix.lines() {
531 let fields: Vec<i32> = line
532 .split_whitespace()
533 .map(|f| f.parse::<i32>().unwrap())
534 .collect();
535 lines.push(fields);
536 }
537 let mut lines_it = lines.into_iter();
538 let header = lines_it.next().unwrap();
539 let forward_size = header[0] as u32;
540 let backward_size = header[1] as u32;
541 let len = 3 + (forward_size * backward_size) as usize;
542 let mut costs = vec![i16::MAX; len];
543 costs[0] = -1;
544 costs[1] = forward_size as i16;
545 costs[2] = backward_size as i16;
546 for fields in lines_it {
547 if fields.is_empty() {
548 continue;
549 }
550 let forward_id = fields[0] as u32;
551 let backward_id = fields[1] as u32;
552 let cost = fields[2] as u16;
553 costs[3 + (forward_id + backward_id * forward_size) as usize] = cost as i16;
554 }
555 costs
556 }
557
558 fn new_costs(matrix: &str) -> Vec<i16> {
560 let bytes = matrix.as_bytes();
561 let header_end = memchr(b'\n', bytes).unwrap_or(bytes.len());
562 let mut header_pos = 0;
563 let forward_size = next_int(&bytes[..header_end], &mut header_pos).unwrap() as u32;
564 let backward_size = next_int(&bytes[..header_end], &mut header_pos).unwrap() as u32;
565 let len = 3 + (forward_size as usize) * (backward_size as usize);
566 let mut costs = vec![i16::MAX; len];
567 costs[0] = -1;
568 costs[1] = forward_size as i16;
569 costs[2] = backward_size as i16;
570 let data = if header_end < bytes.len() {
571 &bytes[header_end + 1..]
572 } else {
573 &[]
574 };
575 fill_costs_sequential(data, forward_size, &mut costs, None).unwrap();
576 costs
577 }
578
579 #[test]
580 fn test_matches_reference_simple() {
581 let matrix = "2 2\n0 0 10\n0 1 20\n1 0 30\n1 1 40\n";
583 assert_eq!(new_costs(matrix), reference_costs(matrix));
584 }
585
586 #[test]
587 fn test_matches_reference_sparse_and_negative() {
588 let matrix = "3 2\n0 0 -1\n2 1 32767\n1 0 -32768\n";
591 let new = new_costs(matrix);
592 let reference = reference_costs(matrix);
593 assert_eq!(new, reference);
594 assert_eq!(new[3], -1);
596 }
597
598 #[test]
599 fn test_no_trailing_newline() {
600 let matrix = "1 1\n0 0 7";
601 assert_eq!(new_costs(matrix), reference_costs(matrix));
602 }
603
604 #[test]
605 fn test_duplicate_last_occurrence_wins() {
606 let matrix = "1 1\n0 0 5\n0 0 9\n";
609 let costs = new_costs(matrix);
610 assert_eq!(costs[3], 9);
611 assert_eq!(costs, reference_costs(matrix));
612 }
613
614 #[cfg(not(target_family = "wasm"))]
615 #[test]
616 fn test_parallel_matches_sequential() {
617 let forward = 200u32;
620 let backward = 200u32;
621 let mut matrix = format!("{forward} {backward}\n");
622 for b in 0..backward {
623 for f in 0..forward {
624 let cost = ((f + b) % 100) as i32 - 50;
625 matrix.push_str(&format!("{f} {b} {cost}\n"));
626 }
627 }
628 let bytes = matrix.as_bytes();
629 let header_end = memchr(b'\n', bytes).unwrap();
630 let data = &bytes[header_end + 1..];
631 let len = 3 + (forward as usize) * (backward as usize);
632
633 let mut seq = vec![i16::MAX; len];
634 seq[0] = -1;
635 seq[1] = forward as i16;
636 seq[2] = backward as i16;
637 fill_costs_sequential(data, forward, &mut seq, None).unwrap();
638
639 let mut par = vec![i16::MAX; len];
640 par[0] = -1;
641 par[1] = forward as i16;
642 par[2] = backward as i16;
643 fill_costs_parallel(data, forward, &mut par, None).unwrap();
644
645 assert_eq!(seq, par);
646 assert_eq!(seq, reference_costs(&matrix));
647 }
648
649 #[test]
650 fn test_missing_field_errors() {
651 let matrix = "2 2\n0 0\n";
653 let bytes = matrix.as_bytes();
654 let header_end = memchr(b'\n', bytes).unwrap();
655 let data = &bytes[header_end + 1..];
656 let mut costs = vec![i16::MAX; 3 + 4];
657 assert!(fill_costs_sequential(data, 2, &mut costs, None).is_err());
658 }
659
660 #[test]
661 fn test_strip_utf8_bom() {
662 let with_bom = [0xEF, 0xBB, 0xBF, b'1', b' ', b'1'];
663 assert_eq!(strip_utf8_bom(&with_bom), b"1 1");
664 assert_eq!(strip_utf8_bom(b"1 1"), b"1 1");
665 }
666}