jellyfish_reader/

matrix.rs

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

/// A rectangular binary matrix over GF(2), used by Jellyfish for hash position computation.
///
/// The matrix has dimensions `rows × cols` where `rows ≤ 64`. It is stored in
/// column-major format: each column is a single `u64` word. Matrix-vector
/// multiplication is done using XOR operations.
///
/// This is used to compute the expected position of a k-mer in the binary/sorted
/// file format, enabling efficient binary search for random access queries.
#[derive(Debug, Clone)]
pub struct RectangularBinaryMatrix {
    /// Number of rows (≤ 64).
    rows: usize,
    /// Number of columns.
    cols: usize,
    /// Column data. None represents an identity matrix.
    columns: Option<Vec<u64>>,
}

impl RectangularBinaryMatrix {
    /// Create a new matrix with the given dimensions and column data.
    pub fn new(rows: usize, cols: usize, columns: Vec<u64>) -> Result<Self> {
        if rows > 64 {
            return Err(Error::InvalidHeader(format!(
                "matrix rows ({rows}) exceeds maximum of 64"
            )));
        }
        if columns.len() != cols {
            return Err(Error::InvalidHeader(format!(
                "expected {cols} columns, got {}",
                columns.len()
            )));
        }
        Ok(Self {
            rows,
            cols,
            columns: Some(columns),
        })
    }

    /// Create an identity matrix.
    pub fn identity(size: usize) -> Self {
        Self {
            rows: size.min(64),
            cols: size,
            columns: None,
        }
    }

    /// Parse a matrix from the JSON header representation.
    ///
    /// In the Jellyfish header, a matrix is stored as an array of integers
    /// representing the column values, preceded by dimension information.
    pub fn from_json(value: &serde_json::Value) -> Result<Self> {
        // Jellyfish stores matrices as: {"r": rows, "c": cols, "columns": [...]}
        // or sometimes just as an array of column values with dimensions inferred
        if value.is_null() {
            return Ok(Self::identity(64));
        }

        if let Some(obj) = value.as_object() {
            let rows = obj
                .get("r")
                .and_then(|v| v.as_u64())
                .ok_or_else(|| Error::MissingField("matrix.r".to_string()))?
                as usize;
            let cols = obj
                .get("c")
                .and_then(|v| v.as_u64())
                .ok_or_else(|| Error::MissingField("matrix.c".to_string()))?
                as usize;

            if let Some(identity) = obj.get("identity") {
                if identity.as_bool().unwrap_or(false) {
                    return Ok(Self::identity(rows.max(cols)));
                }
            }

            let columns = if let Some(col_arr) = obj.get("columns").and_then(|v| v.as_array()) {
                col_arr
                    .iter()
                    .map(|v| {
                        v.as_u64()
                            .or_else(|| v.as_i64().map(|i| i as u64))
                            .ok_or_else(|| {
                                Error::InvalidHeader("invalid matrix column value".to_string())
                            })
                    })
                    .collect::<Result<Vec<u64>>>()?
            } else {
                // Try parsing "data" field (alternative format)
                let data = obj.get("data").and_then(|v| v.as_array()).ok_or_else(|| {
                    Error::MissingField("matrix.columns or matrix.data".to_string())
                })?;
                data.iter()
                    .map(|v| {
                        v.as_u64()
                            .or_else(|| v.as_i64().map(|i| i as u64))
                            .ok_or_else(|| {
                                Error::InvalidHeader("invalid matrix data value".to_string())
                            })
                    })
                    .collect::<Result<Vec<u64>>>()?
            };

            Self::new(rows, cols, columns)
        } else {
            Err(Error::InvalidHeader(
                "matrix must be a JSON object".to_string(),
            ))
        }
    }

    /// Number of rows.
    pub fn rows(&self) -> usize {
        self.rows
    }

    /// Number of columns.
    pub fn cols(&self) -> usize {
        self.cols
    }

    /// Whether this is an identity matrix.
    pub fn is_identity(&self) -> bool {
        self.columns.is_none()
    }

    /// Multiply this matrix by a vector (given as packed u64 words).
    ///
    /// The vector represents the k-mer's packed bits. The result is a `u64`
    /// hash position value.
    pub fn times(&self, vector: &[u64]) -> u64 {
        match &self.columns {
            None => {
                // Identity: just return the low bits
                if vector.is_empty() { 0 } else { vector[0] }
            }
            Some(columns) => {
                let mut result = 0u64;
                let mut col_idx = 0;

                for &word in vector {
                    let mut w = word;
                    let remaining = self.cols - col_idx;
                    let bits_in_word = remaining.min(64);

                    for _ in 0..bits_in_word {
                        if w & 1 != 0 {
                            result ^= columns[col_idx];
                        }
                        w >>= 1;
                        col_idx += 1;
                    }
                }

                result
            }
        }
    }
}

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

    #[test]
    fn test_identity_matrix() {
        let m = RectangularBinaryMatrix::identity(64);
        assert!(m.is_identity());
        assert_eq!(m.times(&[0x12345678]), 0x12345678);
        assert_eq!(m.times(&[0xDEADBEEF]), 0xDEADBEEF);
    }

    #[test]
    fn test_identity_empty() {
        let m = RectangularBinaryMatrix::identity(64);
        assert_eq!(m.times(&[]), 0);
    }

    #[test]
    fn test_zero_matrix() {
        let cols = vec![0u64; 8];
        let m = RectangularBinaryMatrix::new(4, 8, cols).unwrap();
        assert_eq!(m.times(&[0xFF]), 0);
    }

    #[test]
    fn test_simple_matrix() {
        // 2x2 identity-like matrix: col0 = 0b01, col1 = 0b10
        let cols = vec![0b01u64, 0b10u64];
        let m = RectangularBinaryMatrix::new(2, 2, cols).unwrap();
        // vector = 0b11 (both bits set) -> result = col0 XOR col1 = 0b01 ^ 0b10 = 0b11
        assert_eq!(m.times(&[0b11]), 0b11);
        // vector = 0b01 -> result = col0 = 0b01
        assert_eq!(m.times(&[0b01]), 0b01);
        // vector = 0b10 -> result = col1 = 0b10
        assert_eq!(m.times(&[0b10]), 0b10);
    }

    #[test]
    fn test_invalid_dimensions() {
        assert!(RectangularBinaryMatrix::new(65, 4, vec![0; 4]).is_err());
        assert!(RectangularBinaryMatrix::new(4, 4, vec![0; 3]).is_err());
    }

    #[test]
    fn test_from_json_identity() {
        let json = serde_json::json!({
            "r": 32,
            "c": 64,
            "identity": true
        });
        let m = RectangularBinaryMatrix::from_json(&json).unwrap();
        assert!(m.is_identity());
    }

    #[test]
    fn test_from_json_with_columns() {
        let json = serde_json::json!({
            "r": 2,
            "c": 2,
            "columns": [1, 2]
        });
        let m = RectangularBinaryMatrix::from_json(&json).unwrap();
        assert!(!m.is_identity());
        assert_eq!(m.rows(), 2);
        assert_eq!(m.cols(), 2);
    }

    #[test]
    fn test_from_json_null() {
        let json = serde_json::Value::Null;
        let m = RectangularBinaryMatrix::from_json(&json).unwrap();
        assert!(m.is_identity());
    }
}