lindera_dictionary/dictionary/
connection_cost_matrix.rs

1use crate::{LinderaResult, error::LinderaErrorKind, util::Data};
2
3use byteorder::{ByteOrder, LittleEndian};
4use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
5
6#[derive(Clone, Archive, RkyvSerialize, RkyvDeserialize)]
7pub struct ConnectionCostMatrix {
8    /// The connection cost matrix data.
9    /// Previously, this was `Data` (byte array) and costs were read using `LittleEndian::read_i16` at runtime.
10    /// Changed to `Vec<i16>` to enable direct array indexing and avoid deserialization overhead during tokenization.
11    pub costs_data: Vec<i16>,
12    pub backward_size: u32,
13    pub forward_size: u32,
14}
15
16impl ConnectionCostMatrix {
17    /// Load a `ConnectionCostMatrix` from raw binary data.
18    ///
19    /// Supports both the new transposed format (header marker `-1`) and the old format.
20    ///
21    /// # Arguments
22    ///
23    /// * `conn_data` - Raw binary data for the connection cost matrix.
24    ///
25    /// # Returns
26    ///
27    /// A `ConnectionCostMatrix`, or an error if the data is too short or malformed.
28    pub fn load(conn_data: impl Into<Data>) -> LinderaResult<ConnectionCostMatrix> {
29        let conn_data = conn_data.into();
30        if conn_data.len() < 4 {
31            return Err(LinderaErrorKind::Deserialize.with_error(anyhow::anyhow!(
32                "Connection cost matrix data too short: {} bytes",
33                conn_data.len()
34            )));
35        }
36
37        let first_v = LittleEndian::read_i16(&conn_data[0..2]);
38
39        if first_v == -1 {
40            // New format (transposed)
41            if conn_data.len() < 6 {
42                return Err(LinderaErrorKind::Deserialize.with_error(anyhow::anyhow!(
43                    "Connection cost matrix header too short for new format: {} bytes",
44                    conn_data.len()
45                )));
46            }
47            let forward_size = LittleEndian::read_i16(&conn_data[2..4]) as u32;
48            let backward_size = LittleEndian::read_i16(&conn_data[4..6]) as u32;
49            let size = conn_data.len() / 2 - 3;
50            let mut costs_data = vec![0i16; size];
51            LittleEndian::read_i16_into(&conn_data[6..], &mut costs_data);
52
53            Ok(ConnectionCostMatrix {
54                costs_data,
55                backward_size,
56                forward_size,
57            })
58        } else {
59            // Old format
60            let forward_size = first_v as u32;
61            let backward_size = LittleEndian::read_i16(&conn_data[2..4]) as u32;
62            let size = conn_data.len() / 2 - 2;
63            let mut old_costs_data = vec![0i16; size];
64            LittleEndian::read_i16_into(&conn_data[4..], &mut old_costs_data);
65
66            // Transpose to new layout in memory
67            let mut costs_data = vec![0i16; size];
68            for f in 0..forward_size {
69                for b in 0..backward_size {
70                    let old_id = (b + f * backward_size) as usize;
71                    let new_id = (f + b * forward_size) as usize;
72                    costs_data[new_id] = old_costs_data[old_id];
73                }
74            }
75
76            Ok(ConnectionCostMatrix {
77                costs_data,
78                backward_size,
79                forward_size,
80            })
81        }
82    }
83
84    #[inline]
85    pub fn cost(&self, forward_id: u32, backward_id: u32) -> i32 {
86        // Context-id access profiling (feature `ctxfreq`); compiled out by default.
87        #[cfg(feature = "ctxfreq")]
88        crate::builder::context_id_remap::record_access(forward_id, backward_id);
89
90        let cost_id = (forward_id + backward_id * self.forward_size) as usize;
91        self.costs_data[cost_id] as i32
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use byteorder::{LittleEndian, WriteBytesExt};
99
100    #[test]
101    fn test_load_transposed() {
102        let mut data = Vec::new();
103        data.write_i16::<LittleEndian>(-1).unwrap(); // version
104        data.write_i16::<LittleEndian>(2).unwrap(); // forward_size
105        data.write_i16::<LittleEndian>(3).unwrap(); // backward_size
106        // [forward_id + backward_id * forward_size]
107        // [0][0], [1][0], [0][1], [1][1], [0][2], [1][2]
108        data.write_i16::<LittleEndian>(10).unwrap();
109        data.write_i16::<LittleEndian>(11).unwrap();
110        data.write_i16::<LittleEndian>(12).unwrap();
111        data.write_i16::<LittleEndian>(13).unwrap();
112        data.write_i16::<LittleEndian>(14).unwrap();
113        data.write_i16::<LittleEndian>(15).unwrap();
114
115        let matrix = ConnectionCostMatrix::load(data).unwrap();
116        assert_eq!(matrix.forward_size, 2);
117        assert_eq!(matrix.backward_size, 3);
118        assert_eq!(matrix.cost(0, 0), 10);
119        assert_eq!(matrix.cost(1, 0), 11);
120        assert_eq!(matrix.cost(0, 1), 12);
121        assert_eq!(matrix.cost(1, 1), 13);
122        assert_eq!(matrix.cost(0, 2), 14);
123        assert_eq!(matrix.cost(1, 2), 15);
124    }
125
126    #[test]
127    fn test_load_old_format() {
128        let mut data = Vec::new();
129        data.write_i16::<LittleEndian>(2).unwrap(); // forward_size
130        data.write_i16::<LittleEndian>(3).unwrap(); // backward_size
131        // Old layout: [backward_id + forward_id * backward_size]
132        // [0][0], [1][0], [2][0], [0][1], [1][1], [2][1]
133        data.write_i16::<LittleEndian>(10).unwrap();
134        data.write_i16::<LittleEndian>(12).unwrap();
135        data.write_i16::<LittleEndian>(14).unwrap();
136        data.write_i16::<LittleEndian>(11).unwrap();
137        data.write_i16::<LittleEndian>(13).unwrap();
138        data.write_i16::<LittleEndian>(15).unwrap();
139
140        let matrix = ConnectionCostMatrix::load(data).unwrap();
141        assert_eq!(matrix.forward_size, 2);
142        assert_eq!(matrix.backward_size, 3);
143        assert_eq!(matrix.cost(0, 0), 10);
144        assert_eq!(matrix.cost(1, 0), 11);
145        assert_eq!(matrix.cost(0, 1), 12);
146        assert_eq!(matrix.cost(1, 1), 13);
147        assert_eq!(matrix.cost(0, 2), 14);
148        assert_eq!(matrix.cost(1, 2), 15);
149    }
150
151    #[test]
152    fn test_load_data_too_short() {
153        let data: Vec<u8> = vec![0x01, 0x02];
154        let result = ConnectionCostMatrix::load(data);
155        assert!(result.is_err());
156    }
157}
lindera_dictionary/dictionary/connection_cost_matrix.rs

lindera_dictionary/dictionary/
connection_cost_matrix.rs
