mini_myers 0.1.2

/// Transposed query representation for efficient SIMD batch processing.
///
/// This structure stores queries in a transposed format.
/// It also precomputes the `peq` (position-equivalent) bitvectors for each nucleotide type.
///
/// # Examples
///
/// ```rust
/// use mini_myers::TQueries;
///
/// let queries = vec![b"ATG".to_vec(), b"TTG".to_vec()];
/// let transposed = TQueries::new(&queries);
/// assert_eq!(transposed.n_queries, 2);
/// assert_eq!(transposed.query_length, 3);
/// ```
use crate::constant::*;
use wide::{i32x8, u8x32};

#[derive(Debug, Clone)]
pub struct TQueries {
    /// SIMD vectors representing transposed queries, one vector per position per block
    /// Layout: vectors[position][block_idx] where each block contains up to 32 queries
    pub vectors: Vec<Vec<u8x32>>,
    /// Length of each query (all queries must have the same length)
    pub query_length: usize,
    /// Number of queries
    pub n_queries: usize,
    /// Number of blocks of 32 queries needed
    pub n_blocks: usize,
    /// Precomputed peq bitvectors keyed by IUPAC mask (0..=15)
    /// Layout: peq_masks[iupac_mask][query_idx]
    pub peq_masks: [Vec<u32>; IUPAC_MASKS],
    /// Precomputed peq bitvectors for each IUPAC mask
    /// Layout: peqs[iupac_mask][simd_vector_idx] where each SIMD vector contains 8 i32 lanes
    pub peqs: [Vec<i32x8>; IUPAC_MASKS],
}

/// Build peqs vectors from precomputed masks.
#[inline(always)]
fn build_peqs_vectors(
    peq_masks: &[Vec<u32>; IUPAC_MASKS],
    nq: usize,
    vectors_in_block: usize,
) -> [Vec<i32x8>; IUPAC_MASKS] {
    let mut peqs: [Vec<i32x8>; IUPAC_MASKS] =
        std::array::from_fn(|_| Vec::with_capacity(vectors_in_block));

    for v in 0..vectors_in_block {
        let base = v * SIMD_LANES;
        for mask_idx in 0..IUPAC_MASKS {
            let mut lane = [0i32; SIMD_LANES];
            let mask_vec = &peq_masks[mask_idx];
            for (lane_idx, lane_slot) in lane.iter_mut().enumerate() {
                let qi = base + lane_idx;
                if qi < nq {
                    *lane_slot = mask_vec[qi] as i32;
                }
            }
            peqs[mask_idx].push(i32x8::new(lane));
        }
    }

    peqs
}

impl TQueries {
    /// Creates a new `TQueries` structure from a slice of query sequences.
    ///
    /// This method transposes the queries and precomputes the peq bitvectors
    /// for the Myers algorithm. All queries must have the same length (which is <= 32).
    ///
    /// # Arguments
    ///
    /// * `queries` - A slice of byte vectors, where each vector represents one query sequence.
    ///   Each query must contain valid IUPAC nucleotide codes.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use mini_myers::TQueries;
    ///
    /// let queries = vec![b"ATG".to_vec(), b"TTG".to_vec()];
    /// let transposed = TQueries::new(&queries);
    /// ```
    pub fn new(queries: &[Vec<u8>]) -> Self {
        assert!(!queries.is_empty(), "No queries provided");
        let query_length = queries[0].len();
        assert!(
            queries.iter().all(|q| q.len() == query_length),
            "All queries must have the same length"
        );

        let n_queries = queries.len();
        assert!(
            query_length > 0 && query_length <= 32,
            "Query length must be 1..=32"
        );

        // Calculate how many blocks of 32 queries we need
        let n_blocks = n_queries.div_ceil(32);

        // Initialize vector data: [position][block][query_in_block]
        let mut vector_data: Vec<Vec<[u8; 32]>> = vec![vec![[0u8; 32]; n_blocks]; query_length];
        let mut peq_masks: [Vec<u32>; IUPAC_MASKS] = std::array::from_fn(|_| vec![0u32; n_queries]);

        for (qi, q) in queries.iter().enumerate() {
            let block_idx = qi / 32;
            let idx_in_block = qi % 32;

            for (pos, &raw_c) in q.iter().enumerate() {
                let encoded = crate::iupac::get_encoded(raw_c);
                assert!(
                    encoded != INVALID_IUPAC,
                    "Query at index {} contains invalid IUPAC character: {:?}",
                    qi,
                    raw_c as char
                );

                vector_data[pos][block_idx][idx_in_block] = raw_c.to_ascii_uppercase();
                let bit = 1u32 << pos;
                if encoded != 0 {
                    for (mask_idx, mask_vec) in peq_masks.iter_mut().enumerate().skip(1) {
                        let mask = mask_idx as u8;
                        if (encoded & mask) != 0 {
                            mask_vec[qi] |= bit;
                        }
                    }
                }
            }
        }

        // Convert to SIMD vectors
        let vectors: Vec<Vec<u8x32>> = vector_data
            .into_iter()
            .map(|pos_blocks| pos_blocks.into_iter().map(u8x32::new).collect())
            .collect();

        let vectors_in_block = n_queries.div_ceil(SIMD_LANES);
        let peqs = build_peqs_vectors(&peq_masks, n_queries, vectors_in_block);

        Self {
            vectors,
            query_length,
            n_queries,
            n_blocks,
            peq_masks,
            peqs,
        }
    }
}