use std::borrow::Cow;
use std::fs::File;
use std::io::{self, BufRead, Write};
use std::path::Path;
use std::sync::Arc;
use encoding_rs::{Encoding, UTF_16BE, UTF_16LE};
use log::debug;
use memchr::memchr;
use crate::LinderaResult;
use crate::dictionary::context_id_map::ContextIdMap;
use crate::error::LinderaErrorKind;
use crate::util::{read_file, write_data};
const UTF8_BOM: &[u8] = &[0xEF, 0xBB, 0xBF];
#[cfg(not(target_family = "wasm"))]
const PARALLEL_THRESHOLD: usize = 1 << 20;
#[derive(Debug)]
pub struct ConnectionCostMatrixBuilder {
encoding: Cow<'static, str>,
context_id_remap: Option<Arc<ContextIdMap>>,
}
#[derive(Debug, Default)]
pub struct ConnectionCostMatrixBuilderOptions {
encoding: Option<Cow<'static, str>>,
context_id_remap: Option<Arc<ContextIdMap>>,
}
impl ConnectionCostMatrixBuilderOptions {
pub fn encoding(&mut self, value: impl Into<Cow<'static, str>>) -> &mut Self {
self.encoding = Some(value.into());
self
}
pub fn context_id_remap(&mut self, value: Option<Arc<ContextIdMap>>) -> &mut Self {
self.context_id_remap = value;
self
}
pub fn builder(&self) -> ConnectionCostMatrixBuilder {
ConnectionCostMatrixBuilder {
encoding: self.encoding.clone().unwrap_or_else(|| "UTF-8".into()),
context_id_remap: self.context_id_remap.clone(),
}
}
}
impl ConnectionCostMatrixBuilder {
pub fn build(&self, input_dir: &Path, output_dir: &Path) -> LinderaResult<()> {
let matrix_data_path = input_dir.join("matrix.def");
debug!("reading {matrix_data_path:?}");
let buffer = read_file(&matrix_data_path)?;
let decoded = self.decode_if_needed(&buffer)?;
let bytes: &[u8] = match &decoded {
Some(decoded) => decoded.as_bytes(),
None => strip_utf8_bom(&buffer),
};
let header_end = memchr(b'\n', bytes).unwrap_or(bytes.len());
let mut header_pos = 0;
let forward_size = next_int(&bytes[..header_end], &mut header_pos).ok_or_else(|| {
LinderaErrorKind::Content
.with_error(anyhow::anyhow!("matrix.def is missing the size header"))
})? as u32;
let backward_size = next_int(&bytes[..header_end], &mut header_pos).ok_or_else(|| {
LinderaErrorKind::Content.with_error(anyhow::anyhow!(
"matrix.def header is missing backward size"
))
})? as u32;
if let Some(remap) = self.context_id_remap.as_deref()
&& (remap.right.len() != forward_size as usize
|| remap.left.len() != backward_size as usize)
{
return Err(LinderaErrorKind::Content.with_error(anyhow::anyhow!(
"context-id remap size mismatch: remap.right={} vs forward_size={}, remap.left={} vs backward_size={}",
remap.right.len(),
forward_size,
remap.left.len(),
backward_size
)));
}
let len = 3 + (forward_size as usize) * (backward_size as usize);
let mut costs = vec![i16::MAX; len];
costs[0] = -1; costs[1] = forward_size as i16;
costs[2] = backward_size as i16;
let data = if header_end < bytes.len() {
&bytes[header_end + 1..]
} else {
&[]
};
self.fill_costs(data, forward_size, &mut costs)?;
let mut matrix_mtx_buffer = Vec::with_capacity(costs.len() * 2);
for cost in &costs {
matrix_mtx_buffer.extend_from_slice(&cost.to_le_bytes());
}
let wtr_matrix_mtx_path = output_dir.join(Path::new("matrix.mtx"));
let mut wtr_matrix_mtx = io::BufWriter::new(
File::create(wtr_matrix_mtx_path)
.map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?,
);
write_data(&matrix_mtx_buffer, &mut wtr_matrix_mtx)?;
wtr_matrix_mtx
.flush()
.map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?;
Ok(())
}
fn decode_if_needed(&self, buffer: &[u8]) -> LinderaResult<Option<String>> {
let encoding =
Encoding::for_label_no_replacement(self.encoding.as_bytes()).ok_or_else(|| {
LinderaErrorKind::Decode
.with_error(anyhow::anyhow!("Invalid encoding: {}", self.encoding))
})?;
let is_utf16 = encoding == UTF_16LE || encoding == UTF_16BE || has_utf16_bom(buffer);
if is_utf16 {
Ok(Some(encoding.decode(buffer).0.into_owned()))
} else {
Ok(None)
}
}
#[cfg(not(target_family = "wasm"))]
fn fill_costs(&self, data: &[u8], forward_size: u32, costs: &mut [i16]) -> LinderaResult<()> {
let remap = self.context_id_remap.as_deref();
if data.len() >= PARALLEL_THRESHOLD {
fill_costs_parallel(data, forward_size, costs, remap)
} else {
fill_costs_sequential(data, forward_size, costs, remap)
}
}
#[cfg(target_family = "wasm")]
fn fill_costs(&self, data: &[u8], forward_size: u32, costs: &mut [i16]) -> LinderaResult<()> {
fill_costs_sequential(data, forward_size, costs, self.context_id_remap.as_deref())
}
}
fn strip_utf8_bom(buffer: &[u8]) -> &[u8] {
buffer.strip_prefix(UTF8_BOM).unwrap_or(buffer)
}
fn has_utf16_bom(buffer: &[u8]) -> bool {
buffer.starts_with(&[0xFF, 0xFE]) || buffer.starts_with(&[0xFE, 0xFF])
}
pub(crate) fn read_matrix_header(input_dir: &Path, _encoding: &str) -> LinderaResult<(u32, u32)> {
let path = input_dir.join("matrix.def");
let file = File::open(&path).map_err(|err| {
LinderaErrorKind::Io
.with_error(anyhow::anyhow!(err))
.add_context(format!("Failed to open matrix.def: {path:?}"))
})?;
let mut reader = io::BufReader::new(file);
let mut line = Vec::new();
reader.read_until(b'\n', &mut line).map_err(|err| {
LinderaErrorKind::Io
.with_error(anyhow::anyhow!(err))
.add_context("Failed to read matrix.def header line")
})?;
let bytes = strip_utf8_bom(&line);
let mut pos = 0;
let forward_size = next_int(bytes, &mut pos).ok_or_else(|| {
LinderaErrorKind::Content
.with_error(anyhow::anyhow!("matrix.def is missing the size header"))
})? as u32;
let backward_size = next_int(bytes, &mut pos).ok_or_else(|| {
LinderaErrorKind::Content.with_error(anyhow::anyhow!(
"matrix.def header is missing backward size"
))
})? as u32;
Ok((forward_size, backward_size))
}
fn next_int(bytes: &[u8], pos: &mut usize) -> Option<i32> {
while *pos < bytes.len() && matches!(bytes[*pos], b' ' | b'\t' | b'\r') {
*pos += 1;
}
if *pos >= bytes.len() {
return None;
}
let negative = bytes[*pos] == b'-';
if negative {
*pos += 1;
}
let start = *pos;
let mut value: i32 = 0;
while *pos < bytes.len() && bytes[*pos].is_ascii_digit() {
value = value
.wrapping_mul(10)
.wrapping_add((bytes[*pos] - b'0') as i32);
*pos += 1;
}
if *pos == start {
return None;
}
Some(if negative { -value } else { value })
}
fn parse_data_line(
line: &[u8],
forward_size: u32,
costs_len: usize,
remap: Option<&ContextIdMap>,
) -> LinderaResult<Option<(usize, i16)>> {
let mut pos = 0;
let Some(forward_id) = next_int(line, &mut pos) else {
return Ok(None);
};
let backward_id = next_int(line, &mut pos).ok_or_else(|| {
LinderaErrorKind::Content
.with_error(anyhow::anyhow!("matrix.def line is missing backward id"))
})?;
let cost = next_int(line, &mut pos).ok_or_else(|| {
LinderaErrorKind::Content.with_error(anyhow::anyhow!("matrix.def line is missing cost"))
})?;
let forward_id = forward_id as u32 as usize;
let backward_id = backward_id as u32 as usize;
let (fwd, bwd) = match remap {
Some(m) => {
if forward_id >= m.right.len() || backward_id >= m.left.len() {
return Err(LinderaErrorKind::Content.with_error(anyhow::anyhow!(
"matrix.def entry ({forward_id}, {backward_id}) is out of range"
)));
}
(m.right[forward_id] as usize, m.left[backward_id] as usize)
}
None => (forward_id, backward_id),
};
let index = 3 + fwd + bwd * forward_size as usize;
if index >= costs_len {
return Err(LinderaErrorKind::Content.with_error(anyhow::anyhow!(
"matrix.def entry ({forward_id}, {backward_id}) is out of range"
)));
}
let cost = (cost as u16) as i16;
Ok(Some((index, cost)))
}
fn fill_costs_sequential(
data: &[u8],
forward_size: u32,
costs: &mut [i16],
remap: Option<&ContextIdMap>,
) -> LinderaResult<()> {
let costs_len = costs.len();
let mut pos = 0;
while pos < data.len() {
let line_end = memchr(b'\n', &data[pos..])
.map(|offset| pos + offset)
.unwrap_or(data.len());
if let Some((index, cost)) =
parse_data_line(&data[pos..line_end], forward_size, costs_len, remap)?
{
costs[index] = cost;
}
pos = line_end + 1;
}
Ok(())
}
#[cfg(not(target_family = "wasm"))]
fn fill_costs_parallel(
data: &[u8],
forward_size: u32,
costs: &mut [i16],
remap: Option<&ContextIdMap>,
) -> LinderaResult<()> {
use rayon::prelude::*;
let costs_len = costs.len();
let n_chunks = (rayon::current_num_threads() * 4).max(1);
let mut bounds = Vec::with_capacity(n_chunks + 1);
bounds.push(0usize);
for i in 1..n_chunks {
let target = data.len() * i / n_chunks;
let last = *bounds.last().unwrap_or(&0);
if target <= last {
continue;
}
if let Some(offset) = memchr(b'\n', &data[target..]) {
let boundary = target + offset + 1;
if boundary > last && boundary < data.len() {
bounds.push(boundary);
}
}
}
bounds.push(data.len());
let chunks: Vec<&[u8]> = bounds.windows(2).map(|w| &data[w[0]..w[1]]).collect();
let partials: Vec<Vec<(usize, i16)>> = chunks
.par_iter()
.map(|chunk| parse_chunk(chunk, forward_size, costs_len, remap))
.collect::<LinderaResult<Vec<_>>>()?;
for partial in &partials {
for &(index, cost) in partial {
costs[index] = cost;
}
}
Ok(())
}
#[cfg(not(target_family = "wasm"))]
fn parse_chunk(
chunk: &[u8],
forward_size: u32,
costs_len: usize,
remap: Option<&ContextIdMap>,
) -> LinderaResult<Vec<(usize, i16)>> {
let mut out = Vec::with_capacity(chunk.len() / 8);
let mut pos = 0;
while pos < chunk.len() {
let line_end = memchr(b'\n', &chunk[pos..])
.map(|offset| pos + offset)
.unwrap_or(chunk.len());
if let Some(entry) = parse_data_line(&chunk[pos..line_end], forward_size, costs_len, remap)?
{
out.push(entry);
}
pos = line_end + 1;
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
fn reference_costs(matrix: &str) -> Vec<i16> {
let mut lines = Vec::new();
for line in matrix.lines() {
let fields: Vec<i32> = line
.split_whitespace()
.map(|f| f.parse::<i32>().unwrap())
.collect();
lines.push(fields);
}
let mut lines_it = lines.into_iter();
let header = lines_it.next().unwrap();
let forward_size = header[0] as u32;
let backward_size = header[1] as u32;
let len = 3 + (forward_size * backward_size) as usize;
let mut costs = vec![i16::MAX; len];
costs[0] = -1;
costs[1] = forward_size as i16;
costs[2] = backward_size as i16;
for fields in lines_it {
if fields.is_empty() {
continue;
}
let forward_id = fields[0] as u32;
let backward_id = fields[1] as u32;
let cost = fields[2] as u16;
costs[3 + (forward_id + backward_id * forward_size) as usize] = cost as i16;
}
costs
}
fn new_costs(matrix: &str) -> Vec<i16> {
let bytes = matrix.as_bytes();
let header_end = memchr(b'\n', bytes).unwrap_or(bytes.len());
let mut header_pos = 0;
let forward_size = next_int(&bytes[..header_end], &mut header_pos).unwrap() as u32;
let backward_size = next_int(&bytes[..header_end], &mut header_pos).unwrap() as u32;
let len = 3 + (forward_size as usize) * (backward_size as usize);
let mut costs = vec![i16::MAX; len];
costs[0] = -1;
costs[1] = forward_size as i16;
costs[2] = backward_size as i16;
let data = if header_end < bytes.len() {
&bytes[header_end + 1..]
} else {
&[]
};
fill_costs_sequential(data, forward_size, &mut costs, None).unwrap();
costs
}
#[test]
fn test_matches_reference_simple() {
let matrix = "2 2\n0 0 10\n0 1 20\n1 0 30\n1 1 40\n";
assert_eq!(new_costs(matrix), reference_costs(matrix));
}
#[test]
fn test_matches_reference_sparse_and_negative() {
let matrix = "3 2\n0 0 -1\n2 1 32767\n1 0 -32768\n";
let new = new_costs(matrix);
let reference = reference_costs(matrix);
assert_eq!(new, reference);
assert_eq!(new[3], -1);
}
#[test]
fn test_no_trailing_newline() {
let matrix = "1 1\n0 0 7";
assert_eq!(new_costs(matrix), reference_costs(matrix));
}
#[test]
fn test_duplicate_last_occurrence_wins() {
let matrix = "1 1\n0 0 5\n0 0 9\n";
let costs = new_costs(matrix);
assert_eq!(costs[3], 9);
assert_eq!(costs, reference_costs(matrix));
}
#[cfg(not(target_family = "wasm"))]
#[test]
fn test_parallel_matches_sequential() {
let forward = 200u32;
let backward = 200u32;
let mut matrix = format!("{forward} {backward}\n");
for b in 0..backward {
for f in 0..forward {
let cost = ((f + b) % 100) as i32 - 50;
matrix.push_str(&format!("{f} {b} {cost}\n"));
}
}
let bytes = matrix.as_bytes();
let header_end = memchr(b'\n', bytes).unwrap();
let data = &bytes[header_end + 1..];
let len = 3 + (forward as usize) * (backward as usize);
let mut seq = vec![i16::MAX; len];
seq[0] = -1;
seq[1] = forward as i16;
seq[2] = backward as i16;
fill_costs_sequential(data, forward, &mut seq, None).unwrap();
let mut par = vec![i16::MAX; len];
par[0] = -1;
par[1] = forward as i16;
par[2] = backward as i16;
fill_costs_parallel(data, forward, &mut par, None).unwrap();
assert_eq!(seq, par);
assert_eq!(seq, reference_costs(&matrix));
}
#[test]
fn test_missing_field_errors() {
let matrix = "2 2\n0 0\n";
let bytes = matrix.as_bytes();
let header_end = memchr(b'\n', bytes).unwrap();
let data = &bytes[header_end + 1..];
let mut costs = vec![i16::MAX; 3 + 4];
assert!(fill_costs_sequential(data, 2, &mut costs, None).is_err());
}
#[test]
fn test_strip_utf8_bom() {
let with_bom = [0xEF, 0xBB, 0xBF, b'1', b' ', b'1'];
assert_eq!(strip_utf8_bom(&with_bom), b"1 1");
assert_eq!(strip_utf8_bom(b"1 1"), b"1 1");
}
}