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, KeyOrdinal, 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 key_at_ordinal(&self, ordinal: KeyOrdinal) -> Result<Option<String>> {
let Some((_, entry_index, entries)) = self.key_entry_at_ordinal(ordinal)? else {
return Ok(None);
};
Ok(Some(entries[entry_index].key.clone()))
}
pub fn keys_at_ordinals(&self, ordinals: &[KeyOrdinal]) -> Result<Vec<Option<String>>> {
let mut results = vec![None; ordinals.len()];
let mut requests = Vec::new();
for (input_index, ordinal) in ordinals.iter().copied().enumerate() {
let ordinal_value = ordinal.get();
if ordinal_value >= self.len() {
continue;
}
let block_index =
self.find_key_block_by_ordinal(ordinal_value)?
.ok_or(Error::InvalidFormat(
"key ordinal not covered by key blocks",
))?;
let block = self
.key_index
.blocks
.get(block_index)
.ok_or(Error::InvalidFormat("key block index out of range"))?;
let entry_offset = ordinal_value
.checked_sub(block.entry_start_index)
.ok_or(Error::InvalidFormat("key ordinal underflow"))?;
let entry_index = usize::try_from(entry_offset)
.map_err(|_| Error::InvalidFormat("key ordinal exceeds usize"))?;
requests.push((block_index, entry_index, input_index, ordinal_value));
}
requests.sort_unstable_by_key(|(block_index, _, _, _)| *block_index);
let mut range_start = 0usize;
while range_start < requests.len() {
let block_index = requests[range_start].0;
let mut range_end = range_start + 1;
while range_end < requests.len() && requests[range_end].0 == block_index {
range_end += 1;
}
let entries = self.decode_key_block(block_index)?;
for (_, entry_index, input_index, ordinal_value) in &requests[range_start..range_end] {
let entry = entries.get(*entry_index).ok_or_else(|| {
Error::InvalidData(format!(
"key ordinal {ordinal_value} resolved past decoded block length {}",
entries.len()
))
})?;
results[*input_index] = Some(entry.key.clone());
}
range_start = range_end;
}
Ok(results)
}
pub fn key_entry_at_ordinal(
&self,
ordinal: KeyOrdinal,
) -> Result<Option<(usize, usize, Arc<Vec<DecodedKeyEntry>>)>> {
let ordinal = ordinal.get();
if ordinal >= self.len() {
return Ok(None);
}
let block_index = self
.find_key_block_by_ordinal(ordinal)?
.ok_or(Error::InvalidFormat(
"key ordinal not covered by key blocks",
))?;
let block = self
.key_index
.blocks
.get(block_index)
.ok_or(Error::InvalidFormat("key block index out of range"))?;
let entry_offset = ordinal
.checked_sub(block.entry_start_index)
.ok_or(Error::InvalidFormat("key ordinal underflow"))?;
let entry_index = usize::try_from(entry_offset)
.map_err(|_| Error::InvalidFormat("key ordinal exceeds usize"))?;
let entries = self.decode_key_block(block_index)?;
if entry_index >= entries.len() {
return Err(Error::InvalidData(format!(
"key ordinal {ordinal} resolved past decoded block length {}",
entries.len()
)));
}
Ok(Some((block_index, entry_index, entries)))
}
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 find_key_block_by_ordinal(&self, ordinal: u64) -> Result<Option<usize>> {
let blocks = &self.key_index.blocks;
let mut left = 0usize;
let mut right = blocks.len();
while left < right {
let mid = (left + right) / 2;
let block = &blocks[mid];
let block_end = block
.entry_start_index
.checked_add(block.entry_count)
.ok_or(Error::InvalidFormat("keyword entry count overflow"))?;
if ordinal < block.entry_start_index {
right = mid;
} else if ordinal >= block_end {
left = mid + 1;
} else {
return Ok(Some(mid));
}
}
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
}