lindera-dictionary 5.0.1

pub mod character_definition;
pub mod connection_cost_matrix;
pub mod context_id_map;
pub mod metadata;
pub mod prefix_dictionary;
pub mod schema;
pub mod unknown_dictionary;

use std::fs;
use std::path::Path;
use std::str;
use std::sync::Arc;

use byteorder::{ByteOrder, LittleEndian};
use once_cell::sync::Lazy;
use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
use serde::{Deserialize, Serialize};

use crate::LinderaResult;
use crate::dictionary::character_definition::CharacterDefinition;
use crate::dictionary::connection_cost_matrix::ConnectionCostMatrix;
use crate::dictionary::context_id_map::ContextIdMap;
use crate::dictionary::metadata::Metadata;
use crate::dictionary::prefix_dictionary::PrefixDictionary;
use crate::dictionary::unknown_dictionary::UnknownDictionary;
use crate::error::LinderaErrorKind;
use crate::loader::character_definition::CharacterDefinitionLoader;
use crate::loader::connection_cost_matrix::ConnectionCostMatrixLoader;
use crate::loader::metadata::MetadataLoader;
use crate::loader::prefix_dictionary::PrefixDictionaryLoader;
use crate::loader::unknown_dictionary::UnknownDictionaryLoader;
use crate::util::Data;
use crate::viterbi::WordEntry;

pub static UNK: Lazy<Vec<&str>> = Lazy::new(|| vec!["UNK"]);

/// `prefix_dictionary` and `connection_cost_matrix` are `Arc`-wrapped so that
/// `Dictionary::clone()` is O(1) regardless of load method (embedded, mmap,
/// or plain filesystem read) -- these two components dominate a dictionary's
/// memory footprint (tens to hundreds of MB), and nothing in this codebase
/// mutates them after construction.
#[derive(Clone)]
pub struct Dictionary {
    pub prefix_dictionary: Arc<PrefixDictionary>,
    pub connection_cost_matrix: Arc<ConnectionCostMatrix>,
    pub character_definition: CharacterDefinition,
    pub unknown_dictionary: UnknownDictionary,
    pub metadata: Metadata,
}

impl Dictionary {
    /// Retrieve the detail fields (POS, etc.) for an unknown word entry.
    pub fn unknown_word_details(&self, word_id: usize) -> Vec<&str> {
        match self.unknown_dictionary.word_details(word_id as u32) {
            Some(details) => details,
            None => UNK.to_vec(),
        }
    }

    pub fn word_details(&self, word_id: usize) -> Vec<&str> {
        if 4 * word_id >= self.prefix_dictionary.words_idx_data.len() {
            return vec![];
        }

        let idx: usize = match LittleEndian::read_u32(
            &self.prefix_dictionary.words_idx_data[4 * word_id..][..4],
        )
        .try_into()
        {
            Ok(value) => value,
            Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
        };
        let data = &self.prefix_dictionary.words_data[idx..];
        let joined_details_len: usize = match LittleEndian::read_u32(data).try_into() {
            Ok(value) => value,
            Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
        };
        let joined_details_bytes: &[u8] =
            &self.prefix_dictionary.words_data[idx + 4..idx + 4 + joined_details_len];

        let mut details = Vec::new();
        for bytes in joined_details_bytes.split(|&b| b == 0) {
            let detail = match str::from_utf8(bytes) {
                Ok(s) => s,
                Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
            };
            details.push(detail);
        }
        details
    }

    /// Load dictionary from a directory containing dictionary files
    pub fn load_from_path(dict_path: &Path) -> LinderaResult<Self> {
        Self::load_from_path_with_options(dict_path, false)
    }

    /// Load dictionary from a directory with options
    ///
    /// `use_mmap` (when the `mmap` feature is enabled) routes
    /// `connection_cost_matrix` and `prefix_dictionary` through memory-mapped
    /// reads instead of plain file reads. This does **not** make either
    /// component lazily memory-resident at runtime: `ConnectionCostMatrix`
    /// always eagerly decodes into an owned `Vec<i16>`, and
    /// `PrefixDictionary`'s double-array trie (`da`) is always eagerly
    /// deserialized into owned daachorse structures (only
    /// `PrefixDictionary`'s `vals_data`/`words_idx_data`/`words_data` remain
    /// genuinely mmap-backed and are read lazily at lookup time). `metadata`
    /// and `character_definition` and `unknown_dictionary` are always
    /// plain-read regardless of this flag. In short, `use_mmap` only avoids
    /// the initial file-read syscall/allocation for the two components it
    /// covers — it does not provide OS-level lazy paging for tokenization.
    /// Separately, `Dictionary::clone()` is O(1) regardless of `use_mmap`,
    /// since `prefix_dictionary`/`connection_cost_matrix` are `Arc`-wrapped.
    pub fn load_from_path_with_options(dict_path: &Path, use_mmap: bool) -> LinderaResult<Self> {
        // Verify that the dictionary directory exists
        if !dict_path.exists() {
            return Err(LinderaErrorKind::Io.with_error(anyhow::anyhow!(
                "Dictionary path does not exist: {}",
                dict_path.display()
            )));
        }

        if !dict_path.is_dir() {
            return Err(LinderaErrorKind::Io.with_error(anyhow::anyhow!(
                "Dictionary path is not a directory: {}",
                dict_path.display()
            )));
        }

        // Load each component from the dictionary directory
        let metadata = MetadataLoader::load(dict_path)?;
        let character_definition = CharacterDefinitionLoader::load(dict_path)?;

        let connection_cost_matrix = {
            #[cfg(feature = "mmap")]
            if use_mmap {
                ConnectionCostMatrixLoader::load_mmap(dict_path)?
            } else {
                ConnectionCostMatrixLoader::load(dict_path)?
            }
            #[cfg(not(feature = "mmap"))]
            ConnectionCostMatrixLoader::load(dict_path)?
        };

        let prefix_dictionary = {
            #[cfg(feature = "mmap")]
            if use_mmap {
                PrefixDictionaryLoader::load_mmap(dict_path)?
            } else {
                PrefixDictionaryLoader::load(dict_path)?
            }
            #[cfg(not(feature = "mmap"))]
            PrefixDictionaryLoader::load(dict_path)?
        };

        let unknown_dictionary = UnknownDictionaryLoader::load(dict_path)?;

        Ok(Dictionary {
            prefix_dictionary: Arc::new(prefix_dictionary),
            connection_cost_matrix: Arc::new(connection_cost_matrix),
            character_definition,
            unknown_dictionary,
            metadata,
        })
    }

    /// Save dictionary to a directory
    pub fn save_to_path(&self, dict_path: &Path) -> LinderaResult<()> {
        // Create directory if it doesn't exist
        fs::create_dir_all(dict_path)
            .map_err(|err| LinderaErrorKind::Io.with_error(anyhow::anyhow!(err)))?;

        // For now, we'll implement this as needed
        // This would require implementing save methods for each component
        todo!("Dictionary saving will be implemented when needed")
    }
}

#[derive(Clone, Serialize, Deserialize, Archive, RkyvSerialize, RkyvDeserialize)]

pub struct UserDictionary {
    pub dict: PrefixDictionary,
}

impl UserDictionary {
    /// Relabel this dictionary's context IDs with a system dictionary's permutation.
    ///
    /// User dictionaries are always compiled in the *original* context-ID space, which
    /// keeps a built `.bin` portable across remapped and un-remapped system
    /// dictionaries. When one is attached to a system dictionary built with
    /// `connection_id_mapping`, its `left_id`/`right_id` must be moved into the same
    /// space, or every connection cost it participates in would address the wrong
    /// matrix cell — silently, since the IDs stay in range.
    ///
    /// Entries live in `vals_data` as a flat [`WordEntry::SERIALIZED_LEN`]-byte stride
    /// with `left_id` at offset 6 and `right_id` at offset 8 (little endian), so this
    /// rewrites those two `u16`s in place. IDs outside the permutation are left
    /// untouched, matching the builder's behaviour for malformed IDs.
    ///
    /// # Arguments
    ///
    /// * `map` - The permutation persisted in the system dictionary's metadata.
    pub fn remap_context_ids(&mut self, map: &ContextIdMap) {
        const LEFT_ID_OFFSET: usize = 6;
        const RIGHT_ID_OFFSET: usize = 8;

        let mut vals = self.dict.vals_data.to_vec();
        for entry in vals.chunks_exact_mut(WordEntry::SERIALIZED_LEN) {
            let left = LittleEndian::read_u16(&entry[LEFT_ID_OFFSET..][..2]);
            let right = LittleEndian::read_u16(&entry[RIGHT_ID_OFFSET..][..2]);
            LittleEndian::write_u16(&mut entry[LEFT_ID_OFFSET..][..2], map.map_left(left));
            LittleEndian::write_u16(&mut entry[RIGHT_ID_OFFSET..][..2], map.map_right(right));
        }
        self.dict.vals_data = Data::Vec(vals);
    }

    pub fn load(user_dict_data: &[u8]) -> LinderaResult<UserDictionary> {
        let mut aligned = rkyv::util::AlignedVec::<16>::new();
        aligned.extend_from_slice(user_dict_data);
        rkyv::from_bytes::<UserDictionary, rkyv::rancor::Error>(&aligned).map_err(|err| {
            LinderaErrorKind::Deserialize.with_error(anyhow::anyhow!(err.to_string()))
        })
    }

    pub fn word_details(&self, word_id: usize) -> Vec<&str> {
        if 4 * word_id >= self.dict.words_idx_data.len() {
            return UNK.to_vec(); // return empty vector if conversion fails
        }
        let idx = LittleEndian::read_u32(&self.dict.words_idx_data[4 * word_id..][..4]);
        let data = &self.dict.words_data[idx as usize..];

        // Parse the data in the same format as main Dictionary
        let joined_details_len: usize = match LittleEndian::read_u32(data).try_into() {
            Ok(value) => value,
            Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
        };
        let joined_details_bytes: &[u8] =
            &self.dict.words_data[idx as usize + 4..idx as usize + 4 + joined_details_len];

        let mut details = Vec::new();
        for bytes in joined_details_bytes.split(|&b| b == 0) {
            let detail = match str::from_utf8(bytes) {
                Ok(s) => s,
                Err(_) => return UNK.to_vec(), // return empty vector if conversion fails
            };
            details.push(detail);
        }
        details
    }
}