jellyfish-reader 0.1.0

//! # jellyfish-reader
//!
//! A pure Rust library for reading [Jellyfish](https://github.com/gmarcais/Jellyfish)
//! k-mer counting output files.
//!
//! Jellyfish is a fast, memory-efficient tool for counting k-mers in DNA sequences,
//! widely used in bioinformatics. This crate provides native Rust readers for
//! Jellyfish's binary and text output formats, with no C/C++ dependencies.
//!
//! ## Features
//!
//! - **Sequential reading** of binary/sorted and text/sorted Jellyfish files
//! - **Random-access queries** via memory-mapped I/O with binary search
//! - **K-mer representation** (`MerDna`) with canonical form, reverse complement,
//!   and all standard operations
//! - **String k-mer extraction** matching Jellyfish's `StringMers` interface
//! - **Auto-format detection** from file headers
//!
//! ## Quick Start
//!
//! ```no_run
//! use jellyfish_reader::{ReadMerFile, MerDna, QueryMerFile};
//!
//! // Sequential reading
//! let reader = ReadMerFile::open("output.jf").unwrap();
//! for result in reader {
//!     let (mer, count) = result.unwrap();
//!     println!("{}: {}", mer, count);
//! }
//!
//! // Random access
//! let qf = QueryMerFile::open("output.jf").unwrap();
//! let mer: MerDna = "ACGTACGTACGTACGTACGTACGTA".parse().unwrap();
//! if let Some(count) = qf.get(&mer) {
//!     println!("Count: {}", count);
//! }
//! ```
//!
//! ## K-mer Operations
//!
//! ```
//! use jellyfish_reader::MerDna;
//!
//! let mer: MerDna = "ACGT".parse().unwrap();
//!
//! // Reverse complement
//! let rc = mer.get_reverse_complement();
//! assert_eq!(rc.to_string(), "ACGT"); // ACGT is a palindrome
//!
//! // Canonical form (lexicographically smaller of self and RC)
//! let canonical = mer.get_canonical();
//!
//! // Extract k-mers from a sequence
//! use jellyfish_reader::StringMers;
//! let kmers: Vec<_> = StringMers::new("ACGTACGT", 4)
//!     .map(|m| m.to_string())
//!     .collect();
//! assert_eq!(kmers, vec!["ACGT", "CGTA", "GTAC", "TACG", "ACGT"]);
//! ```

pub mod binary;
pub mod error;
pub mod header;
pub mod matrix;
pub mod mer;
pub mod query;
pub mod string_mers;
pub mod text;

// Re-exports for convenience
pub use binary::BinaryReader;
pub use error::{Error, Result};
pub use header::FileHeader;
pub use matrix::RectangularBinaryMatrix;
pub use mer::MerDna;
pub use query::QueryMerFile;
pub use string_mers::{StringMers, string_canonicals, string_mers};
pub use text::TextReader;

use std::fs::File;
use std::io::BufReader;
use std::path::Path;

/// Unified sequential reader for Jellyfish output files.
///
/// Automatically detects the file format (binary/sorted or text/sorted)
/// from the header and provides a single iterator interface.
///
/// # Examples
///
/// ```no_run
/// use jellyfish_reader::ReadMerFile;
///
/// let reader = ReadMerFile::open("output.jf").unwrap();
/// for result in reader {
///     let (mer, count) = result.unwrap();
///     println!("{}: {}", mer, count);
/// }
/// ```
pub enum ReadMerFile {
    /// Binary format reader.
    Binary(BinaryReader<BufReader<File>>),
    /// Text format reader.
    Text(TextReader<BufReader<File>>),
}

impl ReadMerFile {
    /// Open a Jellyfish file for sequential reading.
    ///
    /// The format is auto-detected from the file header.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let file = File::open(path.as_ref())?;
        let mut reader = BufReader::new(file);
        let header = FileHeader::read(&mut reader)?;

        match header.format() {
            header::FORMAT_BINARY_SORTED => {
                let binary = BinaryReader::new(reader, &header)?;
                Ok(ReadMerFile::Binary(binary))
            }
            header::FORMAT_TEXT_SORTED => {
                let text = TextReader::new(reader);
                Ok(ReadMerFile::Text(text))
            }
            other => Err(Error::UnsupportedFormat(other.to_string())),
        }
    }
}

impl Iterator for ReadMerFile {
    type Item = Result<(MerDna, u64)>;

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            ReadMerFile::Binary(r) => r.next(),
            ReadMerFile::Text(r) => r.next(),
        }
    }
}

#[cfg(test)]
mod integration_tests {
    use super::*;
    use std::io::Write;
    use tempfile::NamedTempFile;

    /// Create a test binary/sorted Jellyfish file.
    fn create_binary_jf(k: usize, val_len: usize, records: &[(&str, u64)]) -> NamedTempFile {
        let key_len_bits = k * 2;
        let key_len_bytes = (key_len_bits + 7) / 8;

        let mut sorted_records: Vec<(MerDna, u64)> = records
            .iter()
            .map(|(s, c)| (s.parse::<MerDna>().unwrap(), *c))
            .collect();
        sorted_records.sort_by(|a, b| a.0.cmp(&b.0));

        let size = sorted_records.len().next_power_of_two().max(2);
        let json = serde_json::json!({
            "format": "binary/sorted",
            "key_len": key_len_bits,
            "val_len": val_len,
            "counter_len": val_len,
            "size": size,
            "canonical": false
        });
        let json_str = serde_json::to_string(&json).unwrap();

        let mut file = NamedTempFile::new().unwrap();
        write!(file, "{:09}", json_str.len()).unwrap();
        file.write_all(json_str.as_bytes()).unwrap();

        for (mer, count) in &sorted_records {
            let words = mer.words();
            let mut bytes_written = 0;
            for &word in words {
                for byte_idx in 0..8 {
                    if bytes_written >= key_len_bytes {
                        break;
                    }
                    file.write_all(&[(word >> (byte_idx * 8)) as u8]).unwrap();
                    bytes_written += 1;
                }
            }
            for i in 0..val_len {
                file.write_all(&[(count >> (i * 8)) as u8]).unwrap();
            }
        }
        file.flush().unwrap();
        file
    }

    /// Create a test text/sorted Jellyfish file.
    fn create_text_jf(k: usize, records: &[(&str, u64)]) -> NamedTempFile {
        let json = serde_json::json!({
            "format": "text/sorted",
            "key_len": k * 2,
            "counter_len": 8
        });
        let json_str = serde_json::to_string(&json).unwrap();

        let mut file = NamedTempFile::new().unwrap();
        write!(file, "{:09}", json_str.len()).unwrap();
        file.write_all(json_str.as_bytes()).unwrap();

        for (seq, count) in records {
            writeln!(file, "{seq} {count}").unwrap();
        }
        file.flush().unwrap();
        file
    }

    #[test]
    fn test_read_binary_auto_detect() {
        let file = create_binary_jf(4, 4, &[("ACGT", 42), ("AAAA", 10)]);
        let reader = ReadMerFile::open(file.path()).unwrap();

        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results.len(), 2);

        let strs: Vec<String> = results.iter().map(|(m, _)| m.to_string()).collect();
        let counts: Vec<u64> = results.iter().map(|(_, c)| *c).collect();

        assert!(strs.contains(&"ACGT".to_string()));
        assert!(strs.contains(&"AAAA".to_string()));

        let acgt_idx = strs.iter().position(|s| s == "ACGT").unwrap();
        let aaaa_idx = strs.iter().position(|s| s == "AAAA").unwrap();
        assert_eq!(counts[acgt_idx], 42);
        assert_eq!(counts[aaaa_idx], 10);
    }

    #[test]
    fn test_read_text_auto_detect() {
        let file = create_text_jf(4, &[("ACGT", 42), ("AAAA", 10)]);
        let reader = ReadMerFile::open(file.path()).unwrap();

        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results.len(), 2);
        assert_eq!(results[0].0.to_string(), "ACGT");
        assert_eq!(results[0].1, 42);
        assert_eq!(results[1].0.to_string(), "AAAA");
        assert_eq!(results[1].1, 10);
    }

    #[test]
    fn test_binary_and_query_consistent() {
        let records = [
            ("AAAA", 10u64),
            ("ACGT", 42),
            ("CCCC", 7),
            ("GGGG", 33),
            ("TTTT", 100),
        ];
        let file = create_binary_jf(4, 4, &records);

        let reader = ReadMerFile::open(file.path()).unwrap();
        let seq_results: Vec<_> = reader.map(|r| r.unwrap()).collect();

        let qf = QueryMerFile::open(file.path()).unwrap();

        for (mer, count) in &seq_results {
            let query_count = qf.get(mer);
            assert_eq!(
                query_count,
                Some(*count),
                "mismatch for {}: sequential={}, query={:?}",
                mer,
                count,
                query_count
            );
        }
    }

    #[test]
    fn test_roundtrip_various_k() {
        for k in [1, 2, 4, 8, 16, 25, 31, 32, 33] {
            let seq: String = "ACGT".chars().cycle().take(k).collect();
            let file = create_binary_jf(k, 4, &[(&seq, 1)]);
            let reader = ReadMerFile::open(file.path()).unwrap();
            let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
            assert_eq!(results.len(), 1, "k={k}");
            assert_eq!(results[0].0.to_string(), seq, "k={k}");
            assert_eq!(results[0].1, 1, "k={k}");
        }
    }

    #[test]
    fn test_empty_binary_file() {
        let file = create_binary_jf(4, 4, &[]);
        let reader = ReadMerFile::open(file.path()).unwrap();
        let results: Vec<_> = reader.collect();
        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_empty_text_file() {
        let file = create_text_jf(4, &[]);
        let reader = ReadMerFile::open(file.path()).unwrap();
        let results: Vec<_> = reader.collect();
        assert_eq!(results.len(), 0);
    }

    #[test]
    fn test_large_counts() {
        let file = create_binary_jf(4, 8, &[("ACGT", u64::MAX)]);
        let reader = ReadMerFile::open(file.path()).unwrap();
        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results[0].1, u64::MAX);
    }

    #[test]
    fn test_many_records() {
        let bases = ['A', 'C', 'G', 'T'];
        let mut records = Vec::new();
        for &a in &bases {
            for &b in &bases {
                for &c in &bases {
                    records.push(format!("{a}{b}{c}"));
                }
            }
        }
        let records_with_counts: Vec<(&str, u64)> = records
            .iter()
            .enumerate()
            .map(|(i, s)| (s.as_str(), (i + 1) as u64))
            .collect();

        let file = create_binary_jf(3, 4, &records_with_counts);

        let reader = ReadMerFile::open(file.path()).unwrap();
        let results: Vec<_> = reader.map(|r| r.unwrap()).collect();
        assert_eq!(results.len(), 64);

        let qf = QueryMerFile::open(file.path()).unwrap();
        for (seq, expected_count) in &records_with_counts {
            let mer: MerDna = seq.parse().unwrap();
            let count = qf.get(&mer);
            assert_eq!(count, Some(*expected_count), "query mismatch for {seq}");
        }
    }
}