Skip to main content

sudachi/dic/
build.rs

1/*
2 *  Copyright (c) 2021-2026 Works Applications Co., Ltd.
3 *
4 *  Licensed under the Apache License, Version 2.0 (the "License");
5 *  you may not use this file except in compliance with the License.
6 *  You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 *   Unless required by applicable law or agreed to in writing, software
11 *  distributed under the License is distributed on an "AS IS" BASIS,
12 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 *  See the License for the specific language governing permissions and
14 *  limitations under the License.
15 */
16
17use std::io::Write;
18use std::path::Path;
19use std::time::{SystemTime, UNIX_EPOCH};
20
21use crate::dic::build::error::{BuildFailure, DicBuildError, DicCompilationCtx};
22use crate::dic::build::index::IndexBuilder;
23use crate::dic::build::lexicon::{LexiconWriter, StringStore};
24use crate::dic::build::report::{DictPartReport, ReportBuilder, Reporter};
25use crate::dic::build::resolve::{BinDictResolver, ChainedResolver, RawDictResolver};
26use crate::dic::build::util::default_signature;
27use crate::dic::description::Block;
28use crate::dic::grammar::Grammar;
29use crate::dic::lexicon_set::LexiconSet;
30use crate::dic::{DescriptionAccess, DictionaryAccess, LexiconAccess, ReferenceIdAccess};
31use crate::error::SudachiResult;
32use crate::plugin::input_text::InputTextPlugin;
33use crate::plugin::oov::OovProviderPlugin;
34use crate::plugin::path_rewrite::PathRewritePlugin;
35
36pub(crate) mod conn;
37pub(crate) mod csv_schema;
38pub mod error;
39pub(crate) mod index;
40pub(crate) mod lexicon;
41pub(crate) mod parse;
42pub(crate) mod pos;
43pub mod report;
44mod resolve;
45#[cfg(test)]
46mod test;
47mod util;
48
49const MAX_POS_IDS: usize = i16::MAX as usize;
50const MAX_DIC_STRING_LEN: usize = i16::MAX as usize;
51const MAX_ARRAY_LEN: usize = i8::MAX as usize;
52const DICT_BLOCK_SIZE: usize = 4096;
53const DESCRIPTION_MAGIC_BYTES: &[u8] = b"SudachiBinaryDic";
54const DESCRIPTION_VERSION: u64 = 1;
55const DEFAULT_USER_REFERENCE: &str = "system.dic";
56
57pub enum DataSource<'a> {
58    File(&'a Path),
59    Data(&'a [u8]),
60}
61
62pub trait AsDataSource<'a> {
63    fn convert(self) -> DataSource<'a>;
64    fn name(&self) -> String;
65}
66
67impl<'a> AsDataSource<'a> for DataSource<'a> {
68    fn convert(self) -> DataSource<'a> {
69        self
70    }
71
72    fn name(&self) -> String {
73        match self {
74            DataSource::File(p) => p.to_str().map(|s| s.to_owned()).unwrap_or_default(),
75            DataSource::Data(d) => format!("memory ({} bytes)", d.len()),
76        }
77    }
78}
79
80impl<'a> AsDataSource<'a> for &'a Path {
81    fn convert(self) -> DataSource<'a> {
82        DataSource::File(self)
83    }
84    fn name(&self) -> String {
85        self.to_str().map(|s| s.to_owned()).unwrap_or_default()
86    }
87}
88
89impl<'a> AsDataSource<'a> for &'a [u8] {
90    fn convert(self) -> DataSource<'a> {
91        DataSource::Data(self)
92    }
93    fn name(&self) -> String {
94        format!("memory ({} bytes)", self.len())
95    }
96}
97
98impl<'a, const N: usize> AsDataSource<'a> for &'a [u8; N] {
99    fn convert(self) -> DataSource<'a> {
100        DataSource::Data(&self[..])
101    }
102    fn name(&self) -> String {
103        format!("memory ({} bytes)", self.len())
104    }
105}
106
107pub enum NoDic {}
108
109#[derive(Copy, Clone, Eq, PartialEq)]
110enum BuilderStage {
111    Grammar,
112    Lexicon,
113    Resolved,
114}
115
116impl LexiconAccess for NoDic {
117    fn lexicon(&self) -> &LexiconSet<'_> {
118        panic!("there is no lexicon here")
119    }
120}
121
122impl DictionaryAccess for NoDic {
123    fn grammar(&self) -> &Grammar<'_> {
124        panic!("there is no grammar here")
125    }
126
127    fn input_text_plugins(&self) -> &[Box<dyn InputTextPlugin + Sync + Send>] {
128        &[]
129    }
130
131    fn oov_provider_plugins(&self) -> &[Box<dyn OovProviderPlugin + Sync + Send>] {
132        &[]
133    }
134
135    fn path_rewrite_plugins(&self) -> &[Box<dyn PathRewritePlugin + Sync + Send>] {
136        &[]
137    }
138}
139
140impl ReferenceIdAccess for NoDic {
141    fn reference_ids(&self) -> std::collections::HashMap<u32, String> {
142        std::collections::HashMap::new()
143    }
144}
145
146/// Builds a binary dictionary from csv lexicon and connection matrix (optional)
147pub struct DictBuilder<D> {
148    user: bool,
149    lexicon: lexicon::LexiconReader,
150    conn: conn::ConnBuffer,
151    ctx: DicCompilationCtx,
152    compile_time: SystemTime,
153    description: String,
154    signature: String,
155    reference: String,
156    stage: BuilderStage,
157    prebuilt: Option<D>,
158    reporter: Reporter,
159}
160
161impl DictBuilder<NoDic> {
162    /// Creates a new builder for system dictionary
163    pub fn new_system() -> Self {
164        Self::new_empty()
165    }
166}
167
168impl<D: DictionaryAccess + ReferenceIdAccess> DictBuilder<D> {
169    fn new_empty() -> Self {
170        Self {
171            user: false,
172            lexicon: lexicon::LexiconReader::new(),
173            conn: conn::ConnBuffer::new(),
174            ctx: DicCompilationCtx::default(),
175            compile_time: SystemTime::now(),
176            description: String::new(),
177            signature: String::new(),
178            reference: String::new(),
179            stage: BuilderStage::Grammar,
180            prebuilt: None,
181            reporter: Reporter::new(),
182        }
183    }
184}
185
186impl<D: DictionaryAccess + DescriptionAccess + ReferenceIdAccess> DictBuilder<D> {
187    /// Creates a new builder for user dictionary
188    pub fn new_user(system: D) -> Self {
189        let mut bldr = Self::new_empty();
190        bldr.set_user(true);
191        let cm = system.grammar().conn_matrix();
192        bldr.lexicon
193            .set_max_conn_sizes(cm.num_left() as _, cm.num_right() as _);
194        bldr.lexicon.preload_pos(system.grammar());
195        let max_system_entry_id = system
196            .lexicon()
197            .system_word_ids_in_order()
198            .into_iter()
199            .map(|wid| wid.entry().as_raw() as usize)
200            .max()
201            .unwrap_or(usize::MAX);
202        bldr.lexicon.set_max_system_entry_id(max_system_entry_id);
203        let signature = system.description().signature();
204        if !signature.is_empty() {
205            bldr.reference = signature.to_owned();
206        }
207        bldr.prebuilt = Some(system);
208        bldr
209    }
210}
211
212impl<D: DictionaryAccess + ReferenceIdAccess> DictBuilder<D> {
213    /// Set the dictionary compile time to the specified time instead of current time
214    pub fn set_compile_time<T: Into<std::time::SystemTime>>(
215        &mut self,
216        time: T,
217    ) -> std::time::SystemTime {
218        std::mem::replace(&mut self.compile_time, time.into())
219    }
220
221    /// Set the dictionary description
222    pub fn set_description<T: Into<String>>(&mut self, description: T) {
223        self.description = description.into()
224    }
225
226    /// Read the connection matrix from either a file or an in-memory buffer
227    ///
228    /// This API is intended for system dictionary builds.
229    pub fn read_conn<'a, T: AsDataSource<'a> + 'a>(&mut self, data: T) -> SudachiResult<()> {
230        self.ensure_grammar_stage(
231            "read_conn() must be called before reading lexicon or resolving",
232        )?;
233        let report = ReportBuilder::new(data.name()).read();
234        match data.convert() {
235            DataSource::File(p) => self.conn.read_file(p),
236            DataSource::Data(d) => self.conn.read(d),
237        }?;
238        self.lexicon
239            .set_max_conn_sizes(self.conn.left(), self.conn.right());
240        self.reporter.collect(
241            self.conn.left() as usize * self.conn.right() as usize,
242            report,
243        );
244        Ok(())
245    }
246
247    /// Read POS table csv from either a file or an in-memory buffer.
248    ///
249    /// This API is intended for system dictionary builds.
250    pub fn read_pos<'a, T: AsDataSource<'a> + 'a>(&mut self, data: T) -> SudachiResult<usize> {
251        if self.user {
252            return self.ctx.err(BuildFailure::InvalidSplit(
253                "read_pos is not available for user dictionary".to_owned(),
254            ));
255        }
256        self.ensure_grammar_stage("read_pos() must be called before reading lexicon or resolving")?;
257
258        let report = ReportBuilder::new(data.name()).read();
259        let result = match data.convert() {
260            DataSource::File(p) => self.lexicon.read_pos_file(p),
261            DataSource::Data(d) => self.lexicon.read_pos_bytes(d),
262        };
263        self.reporter.collect_r(result, report)
264    }
265
266    /// Read the csv lexicon from either a file or an in-memory buffer
267    pub fn read_lexicon<'a, T: AsDataSource<'a> + 'a>(&mut self, data: T) -> SudachiResult<usize> {
268        self.ensure_lexicon_stage()?;
269        let report = ReportBuilder::new(data.name()).read();
270        let result = match data.convert() {
271            DataSource::File(p) => self.lexicon.read_file(p),
272            DataSource::Data(d) => self.lexicon.read_bytes(d),
273        };
274        let result = self.reporter.collect_r(result, report);
275        if result.is_ok() {
276            self.stage = BuilderStage::Lexicon;
277        }
278        result
279    }
280
281    /// Resolve the dictionary references.
282    ///
283    /// Returns the number of resolved entries
284    pub fn resolve(&mut self) -> SudachiResult<usize> {
285        self.ensure_resolve_stage()?;
286        self.resolve_impl()
287    }
288
289    /// Compile the binary dictionary and write it to the specified sink
290    pub fn compile<W: Write>(&mut self, w: &mut W) -> SudachiResult<()> {
291        self.prepare_description_fields();
292        self.ensure_compile_stage()?;
293
294        let mut buffer = vec![0u8; DICT_BLOCK_SIZE];
295        let mut blocks: Vec<BlockInfo> = Vec::with_capacity(7);
296
297        if !self.user {
298            self.align_to_block(&mut buffer);
299            let start = buffer.len();
300            let report = ReportBuilder::new("conn_matrix");
301            let size = self.conn.write_to(&mut buffer)?;
302            self.reporter.collect(size, report);
303            blocks.push(BlockInfo::new(Block::ConnectionMatrix, start, size));
304        }
305
306        self.align_to_block(&mut buffer);
307        let start = buffer.len();
308        let report = ReportBuilder::new("pos_table");
309        let size = self.lexicon.write_pos_table(&mut buffer)?;
310        self.reporter.collect(size, report);
311        blocks.push(BlockInfo::new(Block::POSTable, start, size));
312
313        let (trie, word_id_table) = self.build_index_data()?;
314        let strings = StringStore::from_entries(self.lexicon.resolved_entries())?;
315
316        self.align_to_block(&mut buffer);
317        let start = buffer.len();
318        let report = ReportBuilder::new("word_id table");
319        buffer.write_all(&word_id_table)?;
320        self.reporter.collect(word_id_table.len(), report);
321        blocks.push(BlockInfo::new(
322            Block::WordPointers,
323            start,
324            word_id_table.len(),
325        ));
326
327        self.align_to_block(&mut buffer);
328        let start = buffer.len();
329        let report = ReportBuilder::new("trie");
330        buffer.write_all(&trie)?;
331        self.reporter.collect(trie.len(), report);
332        blocks.push(BlockInfo::new(Block::TRIEIndex, start, trie.len()));
333
334        self.align_to_block(&mut buffer);
335        let start = buffer.len();
336        let report = ReportBuilder::new("strings");
337        let size = strings.write(&mut buffer)?;
338        self.reporter.collect(size, report);
339        blocks.push(BlockInfo::new(Block::Strings, start, size));
340
341        self.align_to_block(&mut buffer);
342        let start = buffer.len();
343        let mut writer = LexiconWriter::new(
344            self.lexicon.resolved_entries(),
345            &strings,
346            self.user,
347            &mut self.reporter,
348        );
349        let size = writer.write(&mut buffer)?;
350        blocks.push(BlockInfo::new(Block::Entries, start, size));
351
352        self.align_to_block(&mut buffer);
353        let start = buffer.len();
354        let report = ReportBuilder::new("reference_id_table");
355        let size = self.write_reference_id_table(&mut buffer)?;
356        self.reporter.collect(size, report);
357        blocks.push(BlockInfo::new(Block::ReferenceIdTable, start, size));
358
359        let runtime_costs = self
360            .lexicon
361            .resolved_entries()
362            .iter()
363            .any(|e| e.cost == i16::MIN);
364        // phantom entries stay serialized for reference resolution,
365        // but they are excluded from the public entry counts in the description metadata.
366        let num_total_entries = self
367            .lexicon
368            .resolved_entries()
369            .iter()
370            .filter(|e| !e.is_phantom())
371            .count() as u32;
372        let num_indexed_entries = self
373            .lexicon
374            .resolved_entries()
375            .iter()
376            .filter(|e| !e.is_phantom() && e.should_index())
377            .count() as u32;
378        let description = self.serialize_description(
379            &blocks,
380            num_indexed_entries,
381            num_total_entries,
382            runtime_costs,
383        )?;
384        buffer[..description.len()].copy_from_slice(&description);
385
386        w.write_all(&buffer)?;
387        Ok(())
388    }
389
390    /// Return dictionary build report
391    pub fn report(&self) -> &[DictPartReport] {
392        self.reporter.reports()
393    }
394}
395
396// private functions
397impl<D: DictionaryAccess + ReferenceIdAccess> DictBuilder<D> {
398    fn set_user(&mut self, user: bool) {
399        if user && self.reference.is_empty() {
400            self.reference = DEFAULT_USER_REFERENCE.to_owned();
401        }
402        if !user {
403            self.reference.clear();
404        }
405        self.user = user;
406    }
407
408    fn make_resolver(&self) -> SudachiResult<RawDictResolver> {
409        let line_to_wref = self.lexicon.row_word_refs(self.user);
410        self.ctx.transform(RawDictResolver::new(
411            self.lexicon.entries(),
412            line_to_wref,
413            self.user,
414        ))
415    }
416
417    fn resolve_impl(&mut self) -> SudachiResult<usize> {
418        let this_resolver = self.make_resolver()?;
419        let report = ReportBuilder::new("resolve");
420
421        let cnt = match self.prebuilt.as_ref() {
422            Some(d) => {
423                let built_resolver = BinDictResolver::new(d)?;
424                let chained = ChainedResolver::new(this_resolver, built_resolver);
425                self.lexicon.resolve_entries(&chained, self.user)
426            }
427            None => self.lexicon.resolve_entries(&this_resolver, self.user),
428        };
429        let cnt = self.reporter.collect_r(cnt, report);
430        match cnt {
431            Ok(cnt) => {
432                self.stage = BuilderStage::Resolved;
433                Ok(cnt)
434            }
435            Err((split_info, line)) => Err(DicBuildError {
436                file: "<entries>".to_owned(),
437                line,
438                cause: BuildFailure::InvalidSplitWordReference(split_info),
439            }
440            .into()),
441        }
442    }
443
444    /// Set signature.
445    /// System dictionary has a signature string and user dictionary has empty string.
446    fn prepare_description_fields(&mut self) {
447        if self.user {
448            self.signature.clear();
449        } else if self.signature.is_empty() {
450            self.signature = default_signature(self.compile_time, &self.description);
451        }
452    }
453
454    fn ensure_grammar_stage(&self, message: &'static str) -> SudachiResult<()> {
455        if self.stage != BuilderStage::Grammar {
456            return self.ctx.err(BuildFailure::InvalidBuilderState(message));
457        }
458        Ok(())
459    }
460
461    fn ensure_lexicon_stage(&self) -> SudachiResult<()> {
462        if self.stage == BuilderStage::Resolved {
463            return self.ctx.err(BuildFailure::InvalidBuilderState(
464                "read_lexicon() must be called before resolve()",
465            ));
466        }
467        Ok(())
468    }
469
470    fn ensure_resolve_stage(&self) -> SudachiResult<()> {
471        match self.stage {
472            BuilderStage::Grammar => self.ctx.err(BuildFailure::InvalidBuilderState(
473                "resolve() must be called after reading lexicon",
474            )),
475            BuilderStage::Lexicon => Ok(()),
476            BuilderStage::Resolved => self.ctx.err(BuildFailure::InvalidBuilderState(
477                "resolve() cannot be called more than once",
478            )),
479        }
480    }
481
482    fn ensure_compile_stage(&self) -> SudachiResult<()> {
483        if self.stage != BuilderStage::Resolved {
484            return self.ctx.err(BuildFailure::InvalidBuilderState(
485                "compile() must be called after resolve()",
486            ));
487        }
488        Ok(())
489    }
490
491    fn align_to_block(&self, buffer: &mut Vec<u8>) {
492        let rem = buffer.len() % DICT_BLOCK_SIZE;
493        if rem != 0 {
494            buffer.resize(buffer.len() + (DICT_BLOCK_SIZE - rem), 0);
495        }
496    }
497
498    fn build_index_data(&mut self) -> SudachiResult<(Vec<u8>, Vec<u8>)> {
499        let mut index = IndexBuilder::new();
500        let entry_ids = self.lexicon.row_word_ids(0);
501        // Keep non-indexed, non-phantom entries in the word-id table as a
502        // trailing list. This preserves compatibility with the Java
503        // dictionary format, where callers can enumerate all public entries
504        // from WordIdTable even if some of them are intentionally absent from
505        // the trie. Phantom entries stay internal to reference resolution.
506        let mut non_indexed = Vec::new();
507        for (e, wid) in self
508            .lexicon
509            .resolved_entries()
510            .iter()
511            .zip(entry_ids.into_iter())
512        {
513            if e.should_index() {
514                index.add(e.index_form(), wid);
515            } else if !e.is_phantom() {
516                non_indexed.push(wid);
517            }
518        }
519
520        let word_id_table = index.build_word_id_table(&non_indexed)?;
521        let trie = index.build_trie()?;
522        Ok((trie, word_id_table))
523    }
524
525    fn serialize_description(
526        &self,
527        blocks: &[BlockInfo],
528        num_indexed_entries: u32,
529        num_total_entries: u32,
530        runtime_costs: bool,
531    ) -> SudachiResult<Vec<u8>> {
532        let mut out = Vec::with_capacity(DICT_BLOCK_SIZE);
533        out.extend_from_slice(DESCRIPTION_MAGIC_BYTES);
534        out.extend_from_slice(&DESCRIPTION_VERSION.to_le_bytes());
535
536        let secs = self
537            .compile_time
538            .duration_since(UNIX_EPOCH)
539            .map_err(|_| self.ctx.to_sudachi_err(BuildFailure::InvalidCompileTime))?
540            .as_secs();
541        out.extend_from_slice(&secs.to_le_bytes());
542        let flags = if runtime_costs { 1u64 } else { 0u64 };
543        out.extend_from_slice(&flags.to_le_bytes());
544        self.put_utf8_string(&mut out, &self.description)?;
545        self.put_utf8_string(&mut out, &self.signature)?;
546        self.put_utf8_string(&mut out, &self.reference)?;
547        Self::put_varint(&mut out, num_indexed_entries as u64);
548        Self::put_varint(&mut out, num_total_entries as u64);
549        Self::put_varint(&mut out, blocks.len() as u64);
550        for block in blocks {
551            self.put_utf8_string(&mut out, &block.name)?;
552            Self::put_varint(&mut out, block.start as u64);
553            Self::put_varint(&mut out, block.size as u64);
554        }
555        if out.len() > DICT_BLOCK_SIZE {
556            return self.ctx.err(BuildFailure::InvalidSize {
557                actual: out.len(),
558                expected: DICT_BLOCK_SIZE,
559            });
560        }
561        Ok(out)
562    }
563
564    fn put_utf8_string(&self, dst: &mut Vec<u8>, data: &str) -> SudachiResult<()> {
565        let length = u32::try_from(data.len()).map_err(|_| {
566            self.ctx.to_sudachi_err(BuildFailure::InvalidSize {
567                actual: data.len(),
568                expected: u32::MAX as usize,
569            })
570        })?;
571        Self::put_varint(dst, length as u64);
572        dst.extend_from_slice(data.as_bytes());
573        Ok(())
574    }
575
576    fn put_varint(dst: &mut Vec<u8>, mut value: u64) {
577        loop {
578            let mut byte = (value & 0x7f) as u8;
579            value >>= 7;
580            if value != 0 {
581                byte |= 0x80;
582            }
583            dst.push(byte);
584            if value == 0 {
585                break;
586            }
587        }
588    }
589
590    fn write_reference_id_table<W: Write>(&self, dst: &mut W) -> SudachiResult<usize> {
591        let mut out = Vec::new();
592        let mut rows = Vec::new();
593        let mut offset = lexicon::LexiconReader::ENTRY_INITIAL_OFFSET;
594        for entry in self.lexicon.resolved_entries() {
595            let entry_id =
596                (offset >> crate::dic::word_info::WordInfos::WORD_ID_ALIGNMENT_BITS) as u32;
597            if !entry.is_phantom() {
598                if let Some(reference_id) = entry.reference_id() {
599                    rows.push((entry_id, reference_id));
600                }
601            }
602            offset += entry.expected_entry_size();
603        }
604        Self::put_varint(&mut out, rows.len() as u64);
605        for (entry_id, reference_id) in rows {
606            Self::put_varint(&mut out, entry_id as u64);
607            self.put_utf8_string(&mut out, reference_id)?;
608        }
609        dst.write_all(&out)?;
610        Ok(out.len())
611    }
612}
613
614struct BlockInfo {
615    name: String,
616    start: usize,
617    size: usize,
618}
619
620impl BlockInfo {
621    fn new(block: Block, start: usize, size: usize) -> Self {
622        Self {
623            name: block.to_string(),
624            start,
625            size,
626        }
627    }
628}