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: CharacterDefinition,
47    pub unknown_dictionary: UnknownDictionary,
48    pub metadata: 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    pub fn load_from_path(dict_path: &Path) -> LinderaResult<Self> {
94        Self::load_from_path_with_options(dict_path, false)
95    }
96
97    /// Load dictionary from a directory with options
98    ///
99    /// `use_mmap` (when the `mmap` feature is enabled) routes
100    /// `connection_cost_matrix` and `prefix_dictionary` through memory-mapped
101    /// reads instead of plain file reads. This does **not** make either
102    /// component lazily memory-resident at runtime: `ConnectionCostMatrix`
103    /// always eagerly decodes into an owned `Vec<i16>`, and
104    /// `PrefixDictionary`'s double-array trie (`da`) is always eagerly
105    /// deserialized into owned daachorse structures (only
106    /// `PrefixDictionary`'s `vals_data`/`words_idx_data`/`words_data` remain
107    /// genuinely mmap-backed and are read lazily at lookup time). `metadata`
108    /// and `character_definition` and `unknown_dictionary` are always
109    /// plain-read regardless of this flag. In short, `use_mmap` only avoids
110    /// the initial file-read syscall/allocation for the two components it
111    /// covers — it does not provide OS-level lazy paging for tokenization.
112    /// Separately, `Dictionary::clone()` is O(1) regardless of `use_mmap`,
113    /// since `prefix_dictionary`/`connection_cost_matrix` are `Arc`-wrapped.
114    pub fn load_from_path_with_options(dict_path: &Path, use_mmap: bool) -> LinderaResult<Self> {
115        // Verify that the dictionary directory exists
116        if !dict_path.exists() {
117            return Err(LinderaErrorKind::Io.with_error(anyhow::anyhow!(
118                "Dictionary path does not exist: {}",
119                dict_path.display()
120            )));
121        }
122
123        if !dict_path.is_dir() {
124            return Err(LinderaErrorKind::Io.with_error(anyhow::anyhow!(
125                "Dictionary path is not a directory: {}",
126                dict_path.display()
127            )));
128        }
129
130        // Load each component from the dictionary directory
131        let metadata = MetadataLoader::load(dict_path)?;
132        let character_definition = CharacterDefinitionLoader::load(dict_path)?;
133
134        let connection_cost_matrix = {
135            #[cfg(feature = "mmap")]
136            if use_mmap {
137                ConnectionCostMatrixLoader::load_mmap(dict_path)?
138            } else {
139                ConnectionCostMatrixLoader::load(dict_path)?
140            }
141            #[cfg(not(feature = "mmap"))]
142            ConnectionCostMatrixLoader::load(dict_path)?
143        };
144
145        let prefix_dictionary = {
146            #[cfg(feature = "mmap")]
147            if use_mmap {
148                PrefixDictionaryLoader::load_mmap(dict_path)?
149            } else {
150                PrefixDictionaryLoader::load(dict_path)?
151            }
152            #[cfg(not(feature = "mmap"))]
153            PrefixDictionaryLoader::load(dict_path)?
154        };
155
156        let unknown_dictionary = UnknownDictionaryLoader::load(dict_path)?;
157
158        Ok(Dictionary {
159            prefix_dictionary: Arc::new(prefix_dictionary),
160            connection_cost_matrix: Arc::new(connection_cost_matrix),
161            character_definition,
162            unknown_dictionary,
163            metadata,
164        })
165    }
166
167    /// Save dictionary to a directory
168    pub fn save_to_path(&self, dict_path: &Path) -> LinderaResult<()> {
169        // Create directory if it doesn't exist
170        fs::create_dir_all(dict_path)
171            .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?;
172
173        // For now, we'll implement this as needed
174        // This would require implementing save methods for each component
175        todo!("Dictionary saving will be implemented when needed")
176    }
177}
178
179#[derive(Clone, Serialize, Deserialize, Archive, RkyvSerialize, RkyvDeserialize)]
180
181pub struct UserDictionary {
182    pub dict: PrefixDictionary,
183}
184
185impl UserDictionary {
186    /// Relabel this dictionary's context IDs with a system dictionary's permutation.
187    ///
188    /// User dictionaries are always compiled in the *original* context-ID space, which
189    /// keeps a built `.bin` portable across remapped and un-remapped system
190    /// dictionaries. When one is attached to a system dictionary built with
191    /// `connection_id_mapping`, its `left_id`/`right_id` must be moved into the same
192    /// space, or every connection cost it participates in would address the wrong
193    /// matrix cell — silently, since the IDs stay in range.
194    ///
195    /// Entries live in `vals_data` as a flat [`WordEntry::SERIALIZED_LEN`]-byte stride
196    /// with `left_id` at offset 6 and `right_id` at offset 8 (little endian), so this
197    /// rewrites those two `u16`s in place. IDs outside the permutation are left
198    /// untouched, matching the builder's behaviour for malformed IDs.
199    ///
200    /// # Arguments
201    ///
202    /// * `map` - The permutation persisted in the system dictionary's metadata.
203    pub fn remap_context_ids(&mut self, map: &ContextIdMap) {
204        const LEFT_ID_OFFSET: usize = 6;
205        const RIGHT_ID_OFFSET: usize = 8;
206
207        let mut vals = self.dict.vals_data.to_vec();
208        for entry in vals.chunks_exact_mut(WordEntry::SERIALIZED_LEN) {
209            let left = LittleEndian::read_u16(&entry[LEFT_ID_OFFSET..][..2]);
210            let right = LittleEndian::read_u16(&entry[RIGHT_ID_OFFSET..][..2]);
211            LittleEndian::write_u16(&mut entry[LEFT_ID_OFFSET..][..2], map.map_left(left));
212            LittleEndian::write_u16(&mut entry[RIGHT_ID_OFFSET..][..2], map.map_right(right));
213        }
214        self.dict.vals_data = Data::Vec(vals);
215    }
216
217    pub fn load(user_dict_data: &[u8]) -> LinderaResult<UserDictionary> {
218        let mut aligned = rkyv::util::AlignedVec::<16>::new();
219        aligned.extend_from_slice(user_dict_data);
220        rkyv::from_bytes::<UserDictionary, rkyv::rancor::Error>(&aligned).map_err(|err| {
221            LinderaErrorKind::Deserialize.with_error(anyhow::anyhow!(err.to_string()))
222        })
223    }
224
225    pub fn word_details(&self, word_id: usize) -> Vec<&str> {
226        if 4 * word_id >= self.dict.words_idx_data.len() {
227            return UNK.to_vec(); // return empty vector if conversion fails
228        }
229        let idx = LittleEndian::read_u32(&self.dict.words_idx_data[4 * word_id..][..4]);
230        let data = &self.dict.words_data[idx as usize..];
231
232        // Parse the data in the same format as main Dictionary
233        let joined_details_len: usize = match LittleEndian::read_u32(data).try_into() {
234            Ok(value) => value,
235            Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
236        };
237        let joined_details_bytes: &[u8] =
238            &self.dict.words_data[idx as usize + 4..idx as usize + 4 + joined_details_len];
239
240        let mut details = Vec::new();
241        for bytes in joined_details_bytes.split(|&b| b == 0) {
242            let detail = match str::from_utf8(bytes) {
243                Ok(s) => s,
244                Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
245            };
246            details.push(detail);
247        }
248        details
249    }
250}
lindera_dictionary/dictionary.rs

lindera_dictionary/
dictionary.rs
