ruwex 0.1.0

//! Fast page-title lookup for multistream dumps.
//!
//! On first use this builds a compressed FST mapping each page title to the
//! byte offset of the bz2 stream that contains that page, saved next to the
//! dump (derived from the dump's filename) and transparently reused on later
//! runs. A lookup then seeks straight to that stream, decompresses only it
//! (~100 pages), and returns the matching page — turning a by-name lookup
//! from a full linear index scan (tens of seconds on enwiki) into a few
//! milliseconds.
//!
//! Only multistream dumps are supported: the byte offsets in the companion
//! `*-index.txt.bz2` file are what make the dump seekable in the first place.

use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};

use bzip2::read::MultiBzDecoder;
use fst::{Map, MapBuilder};
use memmap2::Mmap;

use crate::dump::multistream;
use crate::dump::xml::DumpParser;
use crate::dump::{Page, SiteInfo};
use crate::error::{Error, Result};

/// Normalizes a page title for lookup: each underscore becomes a space, and
/// leading/trailing whitespace is trimmed. So `"Richard_Dawkins"` and
/// `"  Richard Dawkins "` both normalize to `"Richard Dawkins"`.
///
/// Internal spacing is left exactly as given — runs of spaces are *not*
/// collapsed, since a title could legitimately contain consecutive spaces.
/// First-letter capitalization is likewise left alone: it is wiki- and
/// namespace-dependent (main-namespace titles are usually capitalized, but
/// wikis like Wiktionary are case-sensitive), so imposing it here would break
/// as many lookups as it fixes.
pub fn normalize_title(title: &str) -> String {
    title.replace('_', " ").trim().to_string()
}

/// A memory-mapped title → stream-offset index for one multistream dump.
pub struct TitleIndex {
    map: Map<Mmap>,
    dump_path: PathBuf,
}

impl TitleIndex {
    /// Path of the title index for a dump: the dump's filename with a
    /// trailing `.bz2` replaced by `.title.fst` (else `.title.fst` appended),
    /// as a sibling of the dump.
    pub fn index_path_for(dump_path: &Path) -> PathBuf {
        let name = dump_path
            .file_name()
            .and_then(|n| n.to_str())
            .unwrap_or("dump");
        let base = name.strip_suffix(".bz2").unwrap_or(name);
        dump_path.with_file_name(format!("{base}.title.fst"))
    }

    /// Opens the title index for `dump_path`, building and saving it on first
    /// use (or when the dump is newer than a stale index). Requires a
    /// multistream dump: its `*-index.txt.bz2` file must exist.
    pub fn open_or_build(dump_path: &Path) -> Result<Self> {
        let fst_path = Self::index_path_for(dump_path);
        if !index_is_fresh(&fst_path, dump_path) {
            build(dump_path, &fst_path)?;
        }
        let file = File::open(&fst_path)?;
        // SAFETY: we own the freshly written index file; it is only read.
        let mmap = unsafe { Mmap::map(&file)? };
        let map = Map::new(mmap).map_err(fst_error)?;
        Ok(Self {
            map,
            dump_path: dump_path.to_path_buf(),
        })
    }

    /// Number of indexed titles.
    pub fn len(&self) -> usize {
        self.map.len()
    }

    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    /// Returns the page with this `title`, or `None` if absent. The title is
    /// normalized first ([`normalize_title`]), so `"Richard_Dawkins"` and
    /// `"  Richard   Dawkins "` both find `"Richard Dawkins"`.
    pub fn find_page(&self, title: &str) -> Result<Option<Page>> {
        let title = normalize_title(title);
        let Some(offset) = self.offset(&title) else {
            return Ok(None);
        };
        Ok(self
            .read_block(offset)?
            .into_iter()
            .find(|page| page.title == title))
        // (A title in the index but not in its stream — a corrupt or
        // mismatched index — yields None.)
    }

    /// Byte offset of the bz2 stream containing `title`, from the index.
    pub fn offset(&self, title: &str) -> Option<u64> {
        self.map.get(title.as_bytes())
    }

    /// Decompresses the single bz2 stream at `offset` and parses all of its
    /// pages (~100). Used both for single lookups and for on-demand template
    /// fetching, where callers cache the result per offset.
    pub fn read_block(&self, offset: u64) -> Result<Vec<Page>> {
        let mut file = File::open(&self.dump_path)?;
        let bytes = multistream::read_stream(&mut file, offset)?;
        let mut parser = DumpParser::new(bytes.as_slice());
        let mut pages = Vec::new();
        while let Some(page) = parser.next_page()? {
            pages.push(page);
        }
        Ok(pages)
    }

    /// Reads the dump's `<siteinfo>` header (the first stream), so single-page
    /// output can carry correct document URLs.
    pub fn site_info(&self) -> Result<SiteInfo> {
        let mut file = File::open(&self.dump_path)?;
        let header = multistream::read_stream(&mut file, 0)?;
        Ok(DumpParser::new(header.as_slice())
            .site_info()?
            .cloned()
            .unwrap_or_default())
    }
}

/// True if `fst_path` exists and is at least as new as `dump_path`.
fn index_is_fresh(fst_path: &Path, dump_path: &Path) -> bool {
    let (Ok(fst_meta), Ok(dump_meta)) = (fst_path.metadata(), dump_path.metadata()) else {
        return false; // index missing (or dump missing — build will error clearly)
    };
    match (fst_meta.modified(), dump_meta.modified()) {
        (Ok(fst_time), Ok(dump_time)) => fst_time >= dump_time,
        _ => true, // timestamps unavailable: trust the existing index
    }
}

/// Builds the title FST from the dump's companion `*-index.txt.bz2` and
/// writes it atomically (temp file + rename, so an interrupted build never
/// leaves a half-written index behind).
fn build(dump_path: &Path, fst_path: &Path) -> Result<()> {
    let index_path = multistream::index_path_for(dump_path).ok_or_else(|| {
        Error::TitleIndex(
            "--title requires a multistream dump with its *-index.txt.bz2 file alongside"
                .to_string(),
        )
    })?;
    log::info!(
        "building title index from {} (first use; one-time step)",
        index_path.display()
    );

    let mut entries = read_title_offsets(&index_path)?;
    // FST keys must be inserted unique and in byte-lexicographic order.
    entries.sort_unstable_by(|a, b| a.0.cmp(&b.0));
    entries.dedup_by(|a, b| a.0 == b.0);

    let tmp_path = fst_path.with_extension("fst.tmp");
    let writer = BufWriter::new(File::create(&tmp_path)?);
    let mut builder = MapBuilder::new(writer).map_err(fst_error)?;
    for (title, offset) in &entries {
        builder
            .insert(title.as_bytes(), *offset)
            .map_err(fst_error)?;
    }
    // into_inner finalizes the FST and hands back the writer to flush.
    let mut writer = builder.into_inner().map_err(fst_error)?;
    writer.flush()?;
    drop(writer);
    std::fs::rename(&tmp_path, fst_path)?;

    log::info!(
        "title index built: {} titles -> {}",
        entries.len(),
        fst_path.display()
    );
    Ok(())
}

/// Reads `(title, stream offset)` pairs from a multistream index file
/// (`offset:pageid:title` per line), parsing bytes directly to avoid
/// per-line UTF-8 validation of the numeric fields.
fn read_title_offsets(index_path: &Path) -> Result<Vec<(String, u64)>> {
    let file = File::open(index_path)?;
    let mut reader: Box<dyn BufRead> = if index_path.extension().is_some_and(|e| e == "bz2") {
        Box::new(BufReader::with_capacity(
            256 * 1024,
            MultiBzDecoder::new(BufReader::with_capacity(256 * 1024, file)),
        ))
    } else {
        Box::new(BufReader::with_capacity(256 * 1024, file))
    };

    let mut entries = Vec::new();
    let mut line = Vec::with_capacity(128);
    loop {
        line.clear();
        if reader.read_until(b'\n', &mut line)? == 0 {
            break;
        }
        let bytes = line.strip_suffix(b"\n").unwrap_or(&line);
        let Some((offset, title)) = parse_index_line(bytes) else {
            continue; // skip blank or malformed lines rather than fail the build
        };
        entries.push((title.to_string(), offset));
    }
    if entries.is_empty() {
        return Err(Error::TitleIndex(
            "multistream index has no usable entries".to_string(),
        ));
    }
    Ok(entries)
}

/// Parses one `offset:pageid:title` line into `(offset, title)`. The title
/// may itself contain colons, so only the first two colons are separators.
fn parse_index_line(bytes: &[u8]) -> Option<(u64, &str)> {
    let mut i = 0;
    let mut offset: u64 = 0;
    let mut digits = 0;
    while i < bytes.len() && bytes[i].is_ascii_digit() {
        offset = offset
            .checked_mul(10)?
            .checked_add(u64::from(bytes[i] - b'0'))?;
        i += 1;
        digits += 1;
    }
    if digits == 0 || bytes.get(i) != Some(&b':') {
        return None;
    }
    i += 1; // consume first ':'
    // skip the page id up to the second ':'
    while i < bytes.len() && bytes[i] != b':' {
        i += 1;
    }
    if bytes.get(i) != Some(&b':') {
        return None;
    }
    i += 1; // consume second ':'
    let title = std::str::from_utf8(&bytes[i..]).ok()?;
    (!title.is_empty()).then_some((offset, title))
}

fn fst_error(error: fst::Error) -> Error {
    Error::TitleIndex(error.to_string())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_index_lines_including_colons_in_titles() {
        assert_eq!(
            parse_index_line(b"570:12:Anarchism"),
            Some((570, "Anarchism"))
        );
        assert_eq!(
            parse_index_line(b"600:5:Talk:Foo: bar"),
            Some((600, "Talk:Foo: bar"))
        );
        assert_eq!(parse_index_line(b""), None);
        assert_eq!(parse_index_line(b"notanumber:1:X"), None);
        assert_eq!(parse_index_line(b"570:12:"), None);
    }

    #[test]
    fn normalizes_underscores_and_trims() {
        assert_eq!(normalize_title("Richard_Dawkins"), "Richard Dawkins");
        assert_eq!(normalize_title("  Richard Dawkins "), "Richard Dawkins");
        assert_eq!(normalize_title("_Alpha_"), "Alpha");
        assert_eq!(normalize_title("\tTab Separated\t"), "Tab Separated");
        assert_eq!(normalize_title(""), "");
        assert_eq!(normalize_title("___"), "");
        // internal runs are preserved, not collapsed
        assert_eq!(normalize_title("Foo__Bar"), "Foo  Bar");
        assert_eq!(normalize_title("Foo   Bar"), "Foo   Bar");
        // idempotent
        assert_eq!(normalize_title(&normalize_title("a__b c")), "a  b c");
    }

    #[test]
    fn index_path_derives_from_dump_name() {
        let p = TitleIndex::index_path_for(Path::new("/data/enwiki-multistream.xml.bz2"));
        assert_eq!(p, Path::new("/data/enwiki-multistream.xml.title.fst"));
        let q = TitleIndex::index_path_for(Path::new("dump.xml"));
        assert_eq!(q, Path::new("dump.xml.title.fst"));
    }
}