lindera_dictionary/
dictionary.rs

1pub mod character_definition;
2pub mod connection_cost_matrix;
3pub mod context_id_map;
4pub mod metadata;
5pub mod prefix_dictionary;
6pub mod schema;
7pub mod unknown_dictionary;
8
9use std::fs;
10use std::path::Path;
11use std::str;
12use std::sync::Arc;
13
14use byteorder::{ByteOrder, LittleEndian};
15use once_cell::sync::Lazy;
16use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
17use serde::{Deserialize, Serialize};
18
19use crate::LinderaResult;
20use crate::dictionary::character_definition::CharacterDefinition;
21use crate::dictionary::connection_cost_matrix::ConnectionCostMatrix;
22use crate::dictionary::context_id_map::ContextIdMap;
23use crate::dictionary::metadata::Metadata;
24use crate::dictionary::prefix_dictionary::PrefixDictionary;
25use crate::dictionary::unknown_dictionary::UnknownDictionary;
26use crate::error::LinderaErrorKind;
27use crate::loader::character_definition::CharacterDefinitionLoader;
28use crate::loader::connection_cost_matrix::ConnectionCostMatrixLoader;
29use crate::loader::metadata::MetadataLoader;
30use crate::loader::prefix_dictionary::PrefixDictionaryLoader;
31use crate::loader::unknown_dictionary::UnknownDictionaryLoader;
32use crate::util::Data;
33use crate::viterbi::WordEntry;
34
35pub static UNK: Lazy<Vec<&str>> = Lazy::new(|| vec!["UNK"]);
36
37/// `prefix_dictionary` and `connection_cost_matrix` are `Arc`-wrapped so that
38/// `Dictionary::clone()` is O(1) regardless of load method (embedded, mmap,
39/// or plain filesystem read) -- these two components dominate a dictionary's
40/// memory footprint (tens to hundreds of MB), and nothing in this codebase
41/// mutates them after construction.
42#[derive(Clone)]
43pub struct Dictionary {
44    pub prefix_dictionary: Arc<PrefixDictionary>,
45    pub connection_cost_matrix: Arc<ConnectionCostMatrix>,
46    pub character_definition: Arc<CharacterDefinition>,
47    pub unknown_dictionary: Arc<UnknownDictionary>,
48    pub metadata: Arc<Metadata>,
49}
50
51impl Dictionary {
52    /// Retrieve the detail fields (POS, etc.) for an unknown word entry.
53    pub fn unknown_word_details(&self, word_id: usize) -> Vec<&str> {
54        match self.unknown_dictionary.word_details(word_id as u32) {
55            Some(details) => details,
56            None => UNK.to_vec(),
57        }
58    }
59
60    pub fn word_details(&self, word_id: usize) -> Vec<&str> {
61        if 4 * word_id >= self.prefix_dictionary.words_idx_data.len() {
62            return vec![];
63        }
64
65        let idx: usize = match LittleEndian::read_u32(
66            &self.prefix_dictionary.words_idx_data[4 * word_id..][..4],
67        )
68        .try_into()
69        {
70            Ok(value) => value,
71            Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
72        };
73        let data = &self.prefix_dictionary.words_data[idx..];
74        let joined_details_len: usize = match LittleEndian::read_u32(data).try_into() {
75            Ok(value) => value,
76            Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
77        };
78        let joined_details_bytes: &[u8] =
79            &self.prefix_dictionary.words_data[idx + 4..idx + 4 + joined_details_len];
80
81        let mut details = Vec::new();
82        for bytes in joined_details_bytes.split(|&b| b == 0) {
83            let detail = match str::from_utf8(bytes) {
84                Ok(s) => s,
85                Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
86            };
87            details.push(detail);
88        }
89        details
90    }
91
92    /// Load dictionary from a directory containing dictionary files.
93    ///
94    /// When the `mmap` feature is compiled in, the connection-cost matrix
95    /// and word list are routed through memory-mapped reads by default
96    /// (#879); use [`Dictionary::load_from_path_with_options`] with
97    /// `use_mmap = false` to force eager reads.
98    pub fn load_from_path(dict_path: &Path) -> LinderaResult<Self> {
99        Self::load_from_path_with_options(dict_path, cfg!(feature = "mmap"))
100    }
101
102    /// Load dictionary from a directory with options
103    ///
104    /// `use_mmap` (when the `mmap` feature is enabled) routes
105    /// `connection_cost_matrix` and `prefix_dictionary` through memory-mapped
106    /// reads instead of plain file reads. This does **not** make either
107    /// component lazily memory-resident at runtime: `ConnectionCostMatrix`
108    /// always eagerly decodes into an owned `Vec<i16>`, and
109    /// `PrefixDictionary`'s double-array trie (`da`) is always eagerly
110    /// deserialized into owned daachorse structures (only
111    /// `PrefixDictionary`'s `vals_data`/`words_idx_data`/`words_data` remain
112    /// genuinely mmap-backed and are read lazily at lookup time). `metadata`
113    /// and `character_definition` and `unknown_dictionary` are always
114    /// plain-read regardless of this flag. In short, `use_mmap` only avoids
115    /// the initial file-read syscall/allocation for the two components it
116    /// covers — it does not provide OS-level lazy paging for tokenization.
117    /// Separately, `Dictionary::clone()` is O(1) regardless of `use_mmap`,
118    /// since `prefix_dictionary`/`connection_cost_matrix` are `Arc`-wrapped.
119    pub fn load_from_path_with_options(dict_path: &Path, use_mmap: bool) -> LinderaResult<Self> {
120        // Verify that the dictionary directory exists
121        if !dict_path.exists() {
122            return Err(LinderaErrorKind::Io.with_error(anyhow::anyhow!(
123                "Dictionary path does not exist: {}",
124                dict_path.display()
125            )));
126        }
127
128        if !dict_path.is_dir() {
129            return Err(LinderaErrorKind::Io.with_error(anyhow::anyhow!(
130                "Dictionary path is not a directory: {}",
131                dict_path.display()
132            )));
133        }
134
135        // Load each component from the dictionary directory
136        let metadata = MetadataLoader::load(dict_path)?;
137        let character_definition = CharacterDefinitionLoader::load(dict_path)?;
138
139        let connection_cost_matrix = {
140            #[cfg(feature = "mmap")]
141            if use_mmap {
142                ConnectionCostMatrixLoader::load_mmap(dict_path)?
143            } else {
144                ConnectionCostMatrixLoader::load(dict_path)?
145            }
146            #[cfg(not(feature = "mmap"))]
147            ConnectionCostMatrixLoader::load(dict_path)?
148        };
149
150        let prefix_dictionary = {
151            #[cfg(feature = "mmap")]
152            if use_mmap {
153                PrefixDictionaryLoader::load_mmap(dict_path)?
154            } else {
155                PrefixDictionaryLoader::load(dict_path)?
156            }
157            #[cfg(not(feature = "mmap"))]
158            PrefixDictionaryLoader::load(dict_path)?
159        };
160
161        let unknown_dictionary = UnknownDictionaryLoader::load(dict_path)?;
162
163        Ok(Dictionary {
164            prefix_dictionary: Arc::new(prefix_dictionary),
165            connection_cost_matrix: Arc::new(connection_cost_matrix),
166            character_definition: Arc::new(character_definition),
167            unknown_dictionary: Arc::new(unknown_dictionary),
168            metadata: Arc::new(metadata),
169        })
170    }
171
172    /// Save dictionary to a directory
173    pub fn save_to_path(&self, dict_path: &Path) -> LinderaResult<()> {
174        // Create directory if it doesn't exist
175        fs::create_dir_all(dict_path)
176            .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?;
177
178        // For now, we'll implement this as needed
179        // This would require implementing save methods for each component
180        todo!("Dictionary saving will be implemented when needed")
181    }
182}
183
184#[derive(Clone, Serialize, Deserialize, Archive, RkyvSerialize, RkyvDeserialize)]
185
186pub struct UserDictionary {
187    pub dict: PrefixDictionary,
188}
189
190impl UserDictionary {
191    /// Relabel this dictionary's context IDs with a system dictionary's permutation.
192    ///
193    /// User dictionaries are always compiled in the *original* context-ID space, which
194    /// keeps a built `.bin` portable across remapped and un-remapped system
195    /// dictionaries. When one is attached to a system dictionary built with
196    /// `connection_id_mapping`, its `left_id`/`right_id` must be moved into the same
197    /// space, or every connection cost it participates in would address the wrong
198    /// matrix cell — silently, since the IDs stay in range.
199    ///
200    /// Entries live in `vals_data` as a flat [`WordEntry::SERIALIZED_LEN`]-byte stride
201    /// with `left_id` at offset 6 and `right_id` at offset 8 (little endian), so this
202    /// rewrites those two `u16`s in place. IDs outside the permutation are left
203    /// untouched, matching the builder's behaviour for malformed IDs.
204    ///
205    /// # Arguments
206    ///
207    /// * `map` - The permutation persisted in the system dictionary's metadata.
208    pub fn remap_context_ids(&mut self, map: &ContextIdMap) {
209        const LEFT_ID_OFFSET: usize = 6;
210        const RIGHT_ID_OFFSET: usize = 8;
211
212        let mut vals = self.dict.vals_data.to_vec();
213        for entry in vals.chunks_exact_mut(WordEntry::SERIALIZED_LEN) {
214            let left = LittleEndian::read_u16(&entry[LEFT_ID_OFFSET..][..2]);
215            let right = LittleEndian::read_u16(&entry[RIGHT_ID_OFFSET..][..2]);
216            LittleEndian::write_u16(&mut entry[LEFT_ID_OFFSET..][..2], map.map_left(left));
217            LittleEndian::write_u16(&mut entry[RIGHT_ID_OFFSET..][..2], map.map_right(right));
218        }
219        self.dict.vals_data = Data::Vec(vals);
220    }
221
222    pub fn load(user_dict_data: &[u8]) -> LinderaResult<UserDictionary> {
223        let mut aligned = rkyv::util::AlignedVec::<16>::new();
224        aligned.extend_from_slice(user_dict_data);
225        rkyv::from_bytes::<UserDictionary, rkyv::rancor::Error>(&aligned).map_err(|err| {
226            LinderaErrorKind::Deserialize.with_error(anyhow::anyhow!(err.to_string()))
227        })
228    }
229
230    pub fn word_details(&self, word_id: usize) -> Vec<&str> {
231        if 4 * word_id >= self.dict.words_idx_data.len() {
232            return UNK.to_vec(); // return empty vector if conversion fails
233        }
234        let idx = LittleEndian::read_u32(&self.dict.words_idx_data[4 * word_id..][..4]);
235        let data = &self.dict.words_data[idx as usize..];
236
237        // Parse the data in the same format as main Dictionary
238        let joined_details_len: usize = match LittleEndian::read_u32(data).try_into() {
239            Ok(value) => value,
240            Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
241        };
242        let joined_details_bytes: &[u8] =
243            &self.dict.words_data[idx as usize + 4..idx as usize + 4 + joined_details_len];
244
245        let mut details = Vec::new();
246        for bytes in joined_details_bytes.split(|&b| b == 0) {
247            let detail = match str::from_utf8(bytes) {
248                Ok(s) => s,
249                Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
250            };
251            details.push(detail);
252        }
253        details
254    }
255}
lindera_dictionary/dictionary.rs

lindera_dictionary/
dictionary.rs
