rustalign_simd/

dp_matrix.rs

//! Dynamic Programming Matrix for SIMD-accelerated alignment
//!
//! This module implements the SSEMatrix structure from aligner_swsse.h,
//! which stores DP matrices in a SIMD-friendly layout.

// SIMD wrapper types (unused in this module but available for future use)

/// Matrix quartet indices matching C++ SSEMatrix
///
/// Each matrix element is a quartet of vectors:
/// - E: gap from above (insertion)
/// - F: gap from left (deletion)
/// - H: diagonal (match/mismatch)
/// - TMP: temporary/reserved space
#[repr(usize)]
pub enum QuartetIndex {
    E = 0,
    F = 1,
    H = 2,
    #[allow(clippy::upper_case_acronyms)]
    TMP = 3,
}

/// SIMD-accelerated DP matrix
///
/// This matches the C++ SSEMatrix structure with the following layout:
/// - Elements (cells) are packed into SSE register vectors
/// - Vectors are packed into quartets (E, F, H, TMP)
/// - Quartets are packed into columns
///
/// The layout is optimized for SIMD operations:
/// ```text
/// Column 0:     Column 1:     Column 2:
/// +----+----+  +----+----+  +----+----+
/// | E0 | H0 |  | E1 | H1 |  | E2 | H2 |
/// +----+----+  +----+----+  +----+----+
/// | F0 | TMP|  | F1 | TMP|  | F2 | TMP|
/// +----+----+  +----+----+  +----+----+
/// ```
pub struct DpMatrix {
    /// Number of vectors per cell (always 4 for E, F, H, TMP)
    nvec_per_cell: usize,
    /// Number of vector rows
    nvecrow: usize,
    /// Number of vector columns
    nveccol: usize,
    /// Row stride in vectors
    rowstride: usize,
    /// Column stride in vectors
    colstride: usize,
    /// Matrix buffer
    matbuf: Vec<u8>,
    /// Whether the matrix has been initialized
    inited: bool,
}

impl DpMatrix {
    /// Create a new DP matrix
    pub fn new() -> Self {
        Self {
            nvec_per_cell: 4,
            nvecrow: 0,
            nveccol: 0,
            rowstride: 0,
            colstride: 0,
            matbuf: Vec::new(),
            inited: false,
        }
    }

    /// Initialize the matrix with the given dimensions
    ///
    /// # Arguments
    /// * `nrow` - Number of rows (query characters + 1)
    /// * `ncol` - Number of columns (reference characters + 1)
    /// * `nvecperrow` - Number of vectors per row
    #[allow(clippy::manual_div_ceil)]
    pub fn init(&mut self, nrow: usize, ncol: usize, nvecperrow: usize) {
        // nvecperrow (wperv in C++) is the number of elements per vector
        // Each cell has 4 vectors (E, F, H, TMP)
        // nvecPerCol = (nrow + wperv - 1) / wperv = ceil(nrow / wperv)
        self.nvecrow = (nrow + nvecperrow - 1) / nvecperrow; // nvecPerCol in C++
        self.nveccol = ncol;
        self.nvec_per_cell = 4;
        self.colstride = 4; // nvecPerCell in C++
        self.rowstride = nvecperrow * 4; // nper * nvecPerCell in C++

        // Total vectors needed: (ncol + 1) * nvecPerCell * nvecPerCol
        let total_vecs = (self.nveccol + 1) * self.nvec_per_cell * self.nvecrow;
        let total_bytes = total_vecs * 16; // 16 bytes per SSE register

        self.matbuf = vec![0u8; total_bytes];
        self.inited = true;
    }

    /// Check if the matrix is initialized
    pub fn is_initialized(&self) -> bool {
        self.inited
    }

    /// Clear the matrix
    pub fn clear(&mut self) {
        self.matbuf.fill(0);
    }

    /// Get the matrix buffer as a pointer
    #[inline]
    pub fn ptr(&mut self) -> *mut u8 {
        assert!(self.inited, "Matrix not initialized");
        self.matbuf.as_mut_ptr()
    }

    /// Get the matrix buffer as a slice
    #[inline]
    pub fn as_slice(&self) -> &[u8] {
        assert!(self.inited, "Matrix not initialized");
        &self.matbuf
    }

    /// Get the matrix buffer as a mutable slice
    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        assert!(self.inited, "Matrix not initialized");
        &mut self.matbuf
    }

    /// Get a pointer to the E vector at the given row and column
    ///
    /// This matches the C++ evec() method.
    #[inline]
    pub fn evec_ptr(&self, row: usize, col: usize) -> usize {
        assert!(
            row < self.nvecrow,
            "Row {} out of bounds (max {})",
            row,
            self.nvecrow
        );
        assert!(
            col < self.nveccol,
            "Col {} out of bounds (max {})",
            col,
            self.nveccol
        );
        let elt = row * self.rowstride + col * self.colstride + QuartetIndex::E as usize;
        assert!(
            elt * 16 < self.matbuf.len(),
            "Element index {} out of bounds",
            elt
        );
        elt * 16 // Byte offset
    }

    /// Get a pointer to the F vector at the given row and column
    ///
    /// This matches the C++ fvec() method.
    #[inline]
    pub fn fvec_ptr(&self, row: usize, col: usize) -> usize {
        assert!(row < self.nvecrow, "Row {} out of bounds", row);
        assert!(col < self.nveccol, "Col {} out of bounds", col);
        let elt = row * self.rowstride + col * self.colstride + QuartetIndex::F as usize;
        assert!(elt * 16 < self.matbuf.len());
        elt * 16
    }

    /// Get a pointer to the H vector at the given row and column
    ///
    /// This matches the C++ hvec() method.
    #[inline]
    pub fn hvec_ptr(&self, row: usize, col: usize) -> usize {
        assert!(row < self.nvecrow, "Row {} out of bounds", row);
        assert!(col < self.nveccol, "Col {} out of bounds", col);
        let elt = row * self.rowstride + col * self.colstride + QuartetIndex::H as usize;
        assert!(elt * 16 < self.matbuf.len());
        elt * 16
    }

    /// Get a pointer to the TMP vector at the given row and column
    ///
    /// This matches the C++ tmpvec() method.
    #[inline]
    pub fn tmpvec_ptr(&self, row: usize, col: usize) -> usize {
        assert!(row < self.nvecrow, "Row {} out of bounds", row);
        assert!(col < self.nveccol, "Col {} out of bounds", col);
        let elt = row * self.rowstride + col * self.colstride + QuartetIndex::TMP as usize;
        assert!(elt * 16 < self.matbuf.len());
        elt * 16
    }

    /// Get the number of vector rows
    pub fn nvecrow(&self) -> usize {
        self.nvecrow
    }

    /// Get the number of vector columns
    pub fn nveccol(&self) -> usize {
        self.nveccol
    }

    /// Get the row stride in vectors
    pub fn rowstride(&self) -> usize {
        self.rowstride
    }

    /// Get the column stride in vectors
    pub fn colstride(&self) -> usize {
        self.colstride
    }

    /// Get the total size of the matrix in bytes
    pub fn size_bytes(&self) -> usize {
        self.matbuf.len()
    }
}

impl Default for DpMatrix {
    fn default() -> Self {
        Self::new()
    }
}

/// Helper structure for working with matrix quartets
///
/// This provides convenient access to the E, F, H, and TMP vectors
/// at a specific matrix position.
pub struct MatrixQuartet<'a> {
    matrix: &'a mut DpMatrix,
    row: usize,
    col: usize,
}

impl<'a> MatrixQuartet<'a> {
    /// Create a new quartet accessor
    pub fn new(matrix: &'a mut DpMatrix, row: usize, col: usize) -> Self {
        assert!(row < matrix.nvecrow());
        assert!(col < matrix.nveccol());
        Self { matrix, row, col }
    }

    /// Get the E vector byte offset
    #[inline]
    pub fn e_off(&self) -> usize {
        self.matrix.evec_ptr(self.row, self.col)
    }

    /// Get the F vector byte offset
    #[inline]
    pub fn f_off(&self) -> usize {
        self.matrix.fvec_ptr(self.row, self.col)
    }

    /// Get the H vector byte offset
    #[inline]
    pub fn h_off(&self) -> usize {
        self.matrix.hvec_ptr(self.row, self.col)
    }

    /// Get the TMP vector byte offset
    #[inline]
    pub fn tmp_off(&self) -> usize {
        self.matrix.tmpvec_ptr(self.row, self.col)
    }

    /// Get mutable references to the matrix buffer
    pub fn buffer(&mut self) -> &mut [u8] {
        self.matrix.as_mut_slice()
    }
}

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

    #[test]
    fn test_dp_matrix_new() {
        let matrix = DpMatrix::new();
        assert!(!matrix.is_initialized());
        assert_eq!(matrix.nvecrow(), 0);
        assert_eq!(matrix.nveccol(), 0);
    }

    #[test]
    fn test_dp_matrix_init() {
        let mut matrix = DpMatrix::new();
        matrix.init(10, 20, 2);
        assert!(matrix.is_initialized());
        // nvecrow = ceil(nrow / nvecperrow) = ceil(10/2) = 5
        assert_eq!(matrix.nvecrow(), 5);
        assert_eq!(matrix.nveccol(), 20);
        assert!(matrix.size_bytes() > 0);
    }

    #[test]
    fn test_dp_matrix_clear() {
        let mut matrix = DpMatrix::new();
        matrix.init(10, 20, 2);
        // Write some data
        matrix.as_mut_slice()[0] = 42;
        matrix.clear();
        assert_eq!(matrix.as_slice()[0], 0);
    }

    #[test]
    fn test_dp_matrix_quartet_offsets() {
        let mut matrix = DpMatrix::new();
        matrix.init(10, 20, 2);
        let e_off = matrix.evec_ptr(0, 0);
        let f_off = matrix.fvec_ptr(0, 0);
        let h_off = matrix.hvec_ptr(0, 0);
        let tmp_off = matrix.tmpvec_ptr(0, 0);

        // Offsets should be in order E < F < H < TMP
        assert!(e_off < f_off);
        assert!(f_off < h_off);
        assert!(h_off < tmp_off);

        // Each vector is 16 bytes
        assert_eq!(f_off - e_off, 16);
        assert_eq!(h_off - f_off, 16);
        assert_eq!(tmp_off - h_off, 16);
    }

    #[test]
    fn test_dp_matrix_strides() {
        let mut matrix = DpMatrix::new();
        matrix.init(10, 20, 2);
        // colstride = nvecPerCell = 4
        assert_eq!(matrix.colstride(), 4);
        // rowstride = nper * nvecPerCell = 2 * 4 = 8
        assert_eq!(matrix.rowstride(), 8);
    }

    #[test]
    fn test_matrix_quartet() {
        let mut matrix = DpMatrix::new();
        matrix.init(10, 20, 2);
        let mut quartet = MatrixQuartet::new(&mut matrix, 3, 5);

        // Get buffer and write some data
        let e_off = quartet.e_off();
        {
            let buffer = quartet.buffer();
            buffer[e_off] = 42;
            buffer[e_off + 1] = 43;

            // Verify data is accessible
            assert_eq!(buffer[e_off], 42);
            assert_eq!(buffer[e_off + 1], 43);
        }
    }
}