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    /// Returns the contiguous cost row for a fixed backward (right-context)
85    /// id, so callers relaxing many forward ids against the same backward id
86    /// pay the offset computation and bounds check once.
87    ///
88    /// # Arguments
89    ///
90    /// * `backward_id` - The backward context id selecting the row.
91    ///
92    /// # Returns
93    ///
94    /// A `forward_size`-long slice indexed directly by forward context id.
95    #[inline]
96    pub fn row(&self, backward_id: u32) -> &[i16] {
97        let start = (backward_id * self.forward_size) as usize;
98        &self.costs_data[start..start + self.forward_size as usize]
99    }
100
101    #[inline]
102    pub fn cost(&self, forward_id: u32, backward_id: u32) -> i32 {
103        // Context-id access profiling (feature `ctxfreq`); compiled out by default.
104        #[cfg(feature = "ctxfreq")]
105        crate::builder::context_id_remap::record_access(forward_id, backward_id);
106
107        let cost_id = (forward_id + backward_id * self.forward_size) as usize;
108        self.costs_data[cost_id] as i32
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use byteorder::{LittleEndian, WriteBytesExt};
116
117    #[test]
118    fn test_load_transposed() {
119        let mut data = Vec::new();
120        data.write_i16::<LittleEndian>(-1).unwrap(); // version
121        data.write_i16::<LittleEndian>(2).unwrap(); // forward_size
122        data.write_i16::<LittleEndian>(3).unwrap(); // backward_size
123        // [forward_id + backward_id * forward_size]
124        // [0][0], [1][0], [0][1], [1][1], [0][2], [1][2]
125        data.write_i16::<LittleEndian>(10).unwrap();
126        data.write_i16::<LittleEndian>(11).unwrap();
127        data.write_i16::<LittleEndian>(12).unwrap();
128        data.write_i16::<LittleEndian>(13).unwrap();
129        data.write_i16::<LittleEndian>(14).unwrap();
130        data.write_i16::<LittleEndian>(15).unwrap();
131
132        let matrix = ConnectionCostMatrix::load(data).unwrap();
133        assert_eq!(matrix.forward_size, 2);
134        assert_eq!(matrix.backward_size, 3);
135        assert_eq!(matrix.cost(0, 0), 10);
136        assert_eq!(matrix.cost(1, 0), 11);
137        assert_eq!(matrix.cost(0, 1), 12);
138        assert_eq!(matrix.cost(1, 1), 13);
139        assert_eq!(matrix.cost(0, 2), 14);
140        assert_eq!(matrix.cost(1, 2), 15);
141    }
142
143    #[test]
144    fn test_load_old_format() {
145        let mut data = Vec::new();
146        data.write_i16::<LittleEndian>(2).unwrap(); // forward_size
147        data.write_i16::<LittleEndian>(3).unwrap(); // backward_size
148        // Old layout: [backward_id + forward_id * backward_size]
149        // [0][0], [1][0], [2][0], [0][1], [1][1], [2][1]
150        data.write_i16::<LittleEndian>(10).unwrap();
151        data.write_i16::<LittleEndian>(12).unwrap();
152        data.write_i16::<LittleEndian>(14).unwrap();
153        data.write_i16::<LittleEndian>(11).unwrap();
154        data.write_i16::<LittleEndian>(13).unwrap();
155        data.write_i16::<LittleEndian>(15).unwrap();
156
157        let matrix = ConnectionCostMatrix::load(data).unwrap();
158        assert_eq!(matrix.forward_size, 2);
159        assert_eq!(matrix.backward_size, 3);
160        assert_eq!(matrix.cost(0, 0), 10);
161        assert_eq!(matrix.cost(1, 0), 11);
162        assert_eq!(matrix.cost(0, 1), 12);
163        assert_eq!(matrix.cost(1, 1), 13);
164        assert_eq!(matrix.cost(0, 2), 14);
165        assert_eq!(matrix.cost(1, 2), 15);
166    }
167
168    #[test]
169    fn test_load_data_too_short() {
170        let data: Vec<u8> = vec![0x01, 0x02];
171        let result = ConnectionCostMatrix::load(data);
172        assert!(result.is_err());
173    }
174}
lindera_dictionary/dictionary/connection_cost_matrix.rs

lindera_dictionary/dictionary/
connection_cost_matrix.rs
