ruwex 0.1.0

//! Support for Wikimedia "multistream" dumps.
//!
//! A `*-pages-articles-multistream.xml.bz2` file is a concatenation of
//! independent bz2 streams: a header stream (`<mediawiki><siteinfo>…`),
//! many streams of ~100 `<page>` elements each, and a `</mediawiki>` footer
//! stream. The companion `*-multistream-index.txt.bz2` file lists
//! `offset:pageid:title` lines, where `offset` is the byte position of the
//! stream containing that page. This lets workers seek and decompress
//! streams in parallel.

use std::fs::File;
use std::io::{BufRead, BufReader, Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};

use bzip2::read::{BzDecoder, MultiBzDecoder};

use crate::error::{Error, Result};

/// Derives the index file path for a multistream dump and checks it exists:
/// `X-multistream.xml.bz2` → `X-multistream-index.txt.bz2`.
pub fn index_path_for(dump_path: &Path) -> Option<PathBuf> {
    let name = dump_path.file_name()?.to_str()?;
    let index_name = name.replace("multistream.xml.bz2", "multistream-index.txt.bz2");
    if index_name == name {
        return None;
    }
    let path = dump_path.with_file_name(index_name);
    path.is_file().then_some(path)
}

/// Reads the distinct stream offsets from an index file (bz2 or plain text),
/// sorted ascending. The header stream at offset 0 is not part of the index.
///
/// Enwiki-scale indexes have tens of millions of lines, so this parses the
/// leading digits from a reused byte buffer instead of allocating per line.
pub fn read_stream_offsets(index_path: &Path) -> Result<Vec<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 offsets: Vec<u64> = Vec::new();
    let mut line = Vec::with_capacity(256);
    loop {
        line.clear();
        if reader.read_until(b'\n', &mut line)? == 0 {
            break;
        }
        let mut offset: u64 = 0;
        let mut digits = 0;
        for &byte in &line {
            match byte {
                b'0'..=b'9' => {
                    offset = offset
                        .checked_mul(10)
                        .and_then(|o| o.checked_add(u64::from(byte - b'0')))
                        .ok_or_else(|| bad_index_line(&line))?;
                    digits += 1;
                }
                b':' => break,
                _ => return Err(bad_index_line(&line)),
            }
        }
        if digits == 0 {
            return Err(bad_index_line(&line));
        }
        if offsets.last() != Some(&offset) {
            offsets.push(offset);
        }
    }
    offsets.sort_unstable();
    offsets.dedup();
    if offsets.is_empty() {
        return Err(Error::InvalidDump("empty multistream index".to_string()));
    }
    Ok(offsets)
}

fn bad_index_line(line: &[u8]) -> Error {
    Error::InvalidDump(format!(
        "bad multistream index line: {:?}",
        String::from_utf8_lossy(line)
    ))
}

/// Decompresses the single bz2 stream starting at `offset`.
pub fn read_stream(file: &mut File, offset: u64) -> Result<Vec<u8>> {
    file.seek(SeekFrom::Start(offset))?;
    let mut decoder = BzDecoder::new(BufReader::with_capacity(64 * 1024, &mut *file));
    let mut bytes = Vec::new();
    decoder.read_to_end(&mut bytes)?;
    Ok(bytes)
}

#[cfg(test)]
mod tests {
    use std::io::Write;

    use bzip2::Compression;
    use bzip2::write::BzEncoder;

    use super::*;

    #[test]
    fn derives_index_path_only_when_present() {
        let dir = tempfile::tempdir().unwrap();
        let dump = dir.path().join("x-pages-articles-multistream.xml.bz2");
        assert_eq!(index_path_for(&dump), None, "index file does not exist yet");

        let index = dir
            .path()
            .join("x-pages-articles-multistream-index.txt.bz2");
        std::fs::write(&index, b"").unwrap();
        assert_eq!(index_path_for(&dump), Some(index));

        let plain = dir.path().join("regular.xml.bz2");
        assert_eq!(index_path_for(&plain), None, "not a multistream name");
    }

    #[test]
    fn reads_and_dedups_offsets() {
        let dir = tempfile::tempdir().unwrap();
        let index = dir.path().join("i-multistream-index.txt.bz2");
        let mut enc = BzEncoder::new(Vec::new(), Compression::best());
        enc.write_all(b"600:1:Alpha\n600:2:B:with:colons\n1500:3:Gamma\n")
            .unwrap();
        std::fs::write(&index, enc.finish().unwrap()).unwrap();

        assert_eq!(read_stream_offsets(&index).unwrap(), vec![600, 1500]);
    }

    #[test]
    fn rejects_garbage_index() {
        let dir = tempfile::tempdir().unwrap();
        let index = dir.path().join("bad-index.txt");
        std::fs::write(&index, b"not a number:1:X\n").unwrap();
        assert!(read_stream_offsets(&index).is_err());
    }
}