mdict-rs 0.1.1

Library-first Rust parser for MDict .mdx and .mdd dictionaries
Documentation
use std::path::Path;
use std::sync::{Arc, Mutex};

use crate::compression::decode_block;
use crate::encoding::TextEncoding;
use crate::error::{Error, Result};
use crate::header::parse_header;
use crate::key_index::{KeyIndex, parse_key_index};
use crate::limits::MAX_RECORD_SPAN_BYTES;
use crate::record_index::{RecordIndex, parse_record_index};
use crate::source::FileSource;
use crate::types::{ContainerKind, Header, OpenOptions};

#[derive(Debug, Clone)]
pub struct DecodedKeyEntry {
    pub key: String,
    pub record_start: u64,
}

pub struct MdictFile {
    pub kind: ContainerKind,
    pub source: Arc<FileSource>,
    pub header: Header,
    pub key_encoding: TextEncoding,
    pub key_index: KeyIndex,
    pub record_index: RecordIndex,
    key_block_cache: Mutex<Option<(usize, Arc<Vec<DecodedKeyEntry>>)>>,
    record_block_cache: Mutex<Option<(usize, Arc<Vec<u8>>)>>,
}

impl std::fmt::Debug for MdictFile {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MdictFile")
            .field("kind", &self.kind)
            .field("source", &self.source)
            .field("header", &self.header)
            .field("key_encoding", &self.key_encoding)
            .field("key_index", &self.key_index)
            .field("record_index", &self.record_index)
            .finish_non_exhaustive()
    }
}

impl MdictFile {
    pub fn open(
        path: impl AsRef<Path>,
        kind: ContainerKind,
        options: &OpenOptions,
    ) -> Result<Self> {
        let source = Arc::new(FileSource::open(path)?);
        let header_section = parse_header(&source, kind)?;
        let key_encoding = TextEncoding::for_container(kind, &header_section.header)?;
        let key_index = parse_key_index(
            &source,
            &header_section.header,
            key_encoding,
            header_section.keyword_section_offset,
            options,
        )?;
        let record_index = parse_record_index(
            &source,
            &header_section.header,
            key_index.record_section_offset,
        )?;
        if key_index.num_entries != record_index.num_entries {
            return Err(Error::InvalidData(format!(
                "entry count mismatch between key index ({}) and record index ({})",
                key_index.num_entries, record_index.num_entries
            )));
        }

        Ok(Self {
            kind,
            source,
            header: header_section.header,
            key_encoding,
            key_index,
            record_index,
            key_block_cache: Mutex::new(None),
            record_block_cache: Mutex::new(None),
        })
    }

    pub fn len(&self) -> u64 {
        self.key_index.num_entries
    }

    pub fn normalize_key(&self, key: &str) -> String {
        self.key_encoding
            .normalize_key(key, self.header.key_case_sensitive, self.header.strip_key)
    }

    pub fn key_block_count(&self) -> usize {
        self.key_index.blocks.len()
    }

    pub fn decode_key_block(&self, index: usize) -> Result<Arc<Vec<DecodedKeyEntry>>> {
        if let Some((cached_index, entries)) = self
            .key_block_cache
            .lock()
            .map_err(|_| Error::InvalidFormat("key block cache mutex poisoned"))?
            .clone()
        {
            if cached_index == index {
                return Ok(entries);
            }
        }

        let block = self
            .key_index
            .blocks
            .get(index)
            .ok_or(Error::InvalidFormat("key block index out of range"))?;
        let block_bytes =
            self.source
                .read_exact_at(block.comp_offset, block.comp_size as usize, "key block")?;
        let decoded = decode_block("key block", &block_bytes, block.decomp_size as usize)?;
        let entries = Arc::new(self.parse_key_block_entries(&decoded, block.entry_count)?);

        let mut cache = self
            .key_block_cache
            .lock()
            .map_err(|_| Error::InvalidFormat("key block cache mutex poisoned"))?;
        *cache = Some((index, Arc::clone(&entries)));
        Ok(entries)
    }

    fn parse_key_block_entries(
        &self,
        bytes: &[u8],
        expected_count: u64,
    ) -> Result<Vec<DecodedKeyEntry>> {
        let mut cursor = crate::cursor::Cursor::new(bytes);
        let mut entries = Vec::with_capacity(expected_count as usize);
        while !cursor.is_empty() {
            let record_start = cursor.read_u64_be("key block record offset")?;
            let start = cursor.offset();
            let (key_bytes, next_offset) =
                self.key_encoding
                    .split_terminated(bytes, start, "key block entry key")?;
            let key = self.key_encoding.decode(key_bytes, "key block entry key")?;
            let key_len = next_offset
                .checked_sub(start)
                .ok_or(Error::InvalidFormat("key block cursor underflow"))?;
            cursor.read_bytes(key_len, "key block entry key bytes")?;
            entries.push(DecodedKeyEntry { key, record_start });
        }
        if entries.len() != expected_count as usize {
            return Err(Error::InvalidData(format!(
                "key block entry count mismatch: expected {expected_count}, decoded {}",
                entries.len()
            )));
        }
        Ok(entries)
    }

    pub fn find_entry(
        &self,
        query: &str,
    ) -> Result<Option<(usize, usize, Arc<Vec<DecodedKeyEntry>>)>> {
        let normalized = self.normalize_key(query);
        for block_index in self.candidate_blocks(&normalized) {
            let entries = self.decode_key_block(block_index)?;
            let mut normalized_hits = None;
            for (entry_index, entry) in entries.iter().enumerate() {
                let candidate = self.normalize_key(&entry.key);
                if candidate == normalized {
                    if entry.key == query {
                        return Ok(Some((block_index, entry_index, entries)));
                    }
                    normalized_hits.get_or_insert((block_index, entry_index));
                }
            }
            if let Some((block_index, entry_index)) = normalized_hits {
                return Ok(Some((block_index, entry_index, entries)));
            }
        }
        Ok(None)
    }

    fn candidate_blocks(&self, query: &str) -> Vec<usize> {
        let blocks = &self.key_index.blocks;
        if blocks.is_empty() {
            return Vec::new();
        }

        let mut left = 0usize;
        let mut right = blocks.len();
        while left < right {
            let mid = (left + right) / 2;
            let block = &blocks[mid];
            if query < block.normalized_first.as_str() {
                right = mid;
            } else if query > block.normalized_last.as_str() {
                left = mid + 1;
            } else {
                return neighbors(mid, blocks.len());
            }
        }

        let mut out = Vec::new();
        if left > 0 {
            out.push(left - 1);
        }
        if left < blocks.len() {
            out.push(left);
        }
        out
    }

    pub fn record_end(
        &self,
        block_index: usize,
        entry_index: usize,
        entries: &[DecodedKeyEntry],
    ) -> Result<u64> {
        if let Some(next) = entries.get(entry_index + 1) {
            return Ok(next.record_start);
        }
        if block_index + 1 < self.key_block_count() {
            let next_block = self.decode_key_block(block_index + 1)?;
            let next = next_block
                .first()
                .ok_or(Error::InvalidFormat("empty key block"))?;
            return Ok(next.record_start);
        }
        Ok(self.record_index.total_decompressed_len)
    }

    pub fn read_record_span(&self, start: u64, end: u64) -> Result<Vec<u8>> {
        let len = checked_record_span_len(self, start, end)?;
        let mut out = Vec::with_capacity(len as usize);
        self.visit_record_span(start, end, |chunk| {
            out.extend_from_slice(chunk);
            Ok(())
        })?;
        Ok(out)
    }

    pub fn visit_record_span<F>(&self, start: u64, end: u64, mut visitor: F) -> Result<()>
    where
        F: FnMut(&[u8]) -> Result<()>,
    {
        checked_record_span_len(self, start, end)?;

        let mut cursor = start;
        while cursor < end {
            let block_index = self
                .record_index
                .find_block(cursor)
                .ok_or(Error::InvalidFormat(
                    "record offset not covered by record blocks",
                ))?;
            let block = &self.record_index.blocks[block_index];
            let decoded = self.decode_record_block(block_index)?;
            let local_start = (cursor - block.decomp_offset) as usize;
            let local_end = usize::min(decoded.len(), (end - block.decomp_offset) as usize);
            visitor(&decoded[local_start..local_end])?;
            cursor = block.decomp_offset + local_end as u64;
        }

        Ok(())
    }

    fn decode_record_block(&self, index: usize) -> Result<Arc<Vec<u8>>> {
        if let Some((cached_index, data)) = self
            .record_block_cache
            .lock()
            .map_err(|_| Error::InvalidFormat("record block cache mutex poisoned"))?
            .clone()
        {
            if cached_index == index {
                return Ok(data);
            }
        }

        let block = self
            .record_index
            .blocks
            .get(index)
            .ok_or(Error::InvalidFormat("record block index out of range"))?;
        let block_bytes = self.source.read_exact_at(
            block.comp_offset,
            block.comp_size as usize,
            "record block",
        )?;
        let decoded = Arc::new(decode_block(
            "record block",
            &block_bytes,
            block.decomp_size as usize,
        )?);

        let mut cache = self
            .record_block_cache
            .lock()
            .map_err(|_| Error::InvalidFormat("record block cache mutex poisoned"))?;
        *cache = Some((index, Arc::clone(&decoded)));
        Ok(decoded)
    }
}

fn checked_record_span_len(file: &MdictFile, start: u64, end: u64) -> Result<u64> {
    if end < start {
        return Err(Error::InvalidFormat("record range is inverted"));
    }
    let len = end
        .checked_sub(start)
        .ok_or(Error::InvalidFormat("record range overflow"))?;
    if len as usize > MAX_RECORD_SPAN_BYTES {
        return Err(Error::LimitExceeded {
            limit: "record_span_bytes",
            value: len,
            max: MAX_RECORD_SPAN_BYTES as u64,
        });
    }
    if end > file.record_index.total_decompressed_len {
        return Err(Error::InvalidFormat("record range exceeds record data"));
    }
    Ok(len)
}

fn neighbors(index: usize, len: usize) -> Vec<usize> {
    let mut out = Vec::with_capacity(3);
    if index > 0 {
        out.push(index - 1);
    }
    out.push(index);
    if index + 1 < len {
        out.push(index + 1);
    }
    out
}