rustalign_fmindex/

locus.rs

//! SideLocus - Position within FM-index
//!
//! A SideLocus represents a position within the FM-index, specifically
//! within a "side" (a fixed-size chunk of the BWT).
//!
//! This matches the C++ SideLocus struct in bt2_idx.h for binary compatibility.

use crate::EbwtParams;
use crate::params::OFF_SIZE;
use std::fmt;

/// Position within the FM-index
///
/// This corresponds to the `SideLocus` struct in the C++ code.
/// The FM-index is organized into "sides" of fixed size for
/// efficient cache access.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SideLocus {
    /// Offset of side within ebwt[] (in bytes)
    pub side_byte_off: u64,

    /// Index of side
    pub side_num: u64,

    /// Character offset within side (0 to sideBwtLen-1)
    pub char_off: u32,

    /// Byte within side (0 to sideBwtSz-1)
    pub by: i32,

    /// Bit-pair within byte (0-3)
    pub bp: i32,
}

impl SideLocus {
    /// Size constant: 48 * OFF_SIZE (hard-coded in C++)
    const SIDE_BWT_LEN: u64 = 48 * OFF_SIZE as u64;

    /// Create a new invalid SideLocus (default state)
    pub fn new() -> Self {
        Self {
            side_byte_off: 0,
            side_num: 0,
            char_off: 0,
            by: -1,
            bp: -1,
        }
    }

    /// Initialize from a BWT row position
    ///
    /// # Arguments
    /// * `row` - BWT row position
    /// * `ep` - Ebwt parameters
    /// * `ebwt` - Pointer to ebwt data (for prefetching)
    ///
    /// # Safety
    /// The ebwt pointer must be valid for the lifetime of this SideLocus.
    pub fn init_from_row(&mut self, row: u64, ep: &EbwtParams, ebwt: &[u8]) {
        let side_sz = ep.side_sz as i32;
        // Side length is hard-coded for now; this allows the compiler
        // to do clever things to accelerate / and %.
        let side_num = row / Self::SIDE_BWT_LEN;
        assert!(side_num < ep.num_sides, "side_num out of range");
        let char_off = (row % Self::SIDE_BWT_LEN) as u32;

        self.side_num = side_num;
        self.char_off = char_off;
        let s_byte_off = side_num * side_sz as u64;

        // Prefetch cache lines
        Self::prefetch_from_ebwt(ebwt, s_byte_off as usize);

        self.side_byte_off = s_byte_off;
        assert!(row <= ep.len, "row exceeds reference length");
        assert!(
            self.side_byte_off + side_sz as u64 <= ep.ebwt_tot_sz,
            "side_byte_off exceeds ebwt size"
        );

        // Tons of cache misses on the next line
        self.by = (char_off >> 2) as i32; // byte within side
        assert!(self.by < ep.side_bwt_sz as i32, "by exceeds side_bwt_sz");
        self.bp = (char_off & 3) as i32; // bit-pair within byte
    }

    /// Initialize two SideLocus objects from a top/bot pair
    ///
    /// This is an optimized version that uses the result from one call
    /// to initFromRow to possibly avoid a second call.
    pub fn init_from_top_bot(top: u64, bot: u64, ep: &EbwtParams, ebwt: &[u8]) -> (Self, Self) {
        assert!(bot > top, "bot must be > top");

        let mut ltop = Self::new();
        ltop.init_from_row(top, ep, ebwt);

        let spread = bot - top;
        let char_off_sum = ltop.char_off as u64 + spread;

        let lbot = if char_off_sum < ep.side_bwt_len as u64 {
            // Same side as ltop
            let bchar_off = char_off_sum as u32;
            Self {
                side_byte_off: ltop.side_byte_off,
                side_num: ltop.side_num,
                char_off: bchar_off,
                by: (bchar_off >> 2) as i32,
                bp: (bchar_off & 3) as i32,
            }
        } else {
            // Different side - need to init from row
            let mut locus = Self::new();
            locus.init_from_row(bot, ep, ebwt);
            locus
        };

        (ltop, lbot)
    }

    /// Prefetch cache lines for a given row
    pub fn prefetch_from_row(row: u64, ep: &EbwtParams, ebwt: &[u8]) {
        let side_sz = ep.side_sz as i32;
        let side_num = row / Self::SIDE_BWT_LEN;
        let side_byte_off = (side_num * side_sz as u64) as usize;
        Self::prefetch_from_ebwt(ebwt, side_byte_off);
    }

    /// Prefetch cache lines for this SideLocus
    pub fn prefetch(&self, ebwt: &[u8]) {
        Self::prefetch_from_ebwt(ebwt, self.side_byte_off as usize);
    }

    /// Prefetch from ebwt at given offset
    fn prefetch_from_ebwt(ebwt: &[u8], offset: usize) {
        if offset + 64 <= ebwt.len() {
            // Prefetch first 64-byte cache line
            #[cfg(target_arch = "x86_64")]
            unsafe {
                std::arch::x86_64::_mm_prefetch(
                    ebwt.as_ptr().add(offset) as *const i8,
                    std::arch::x86_64::_MM_HINT_T0,
                );
            }
            // Prefetch second 64-byte cache line
            if offset + 128 <= ebwt.len() {
                #[cfg(target_arch = "x86_64")]
                unsafe {
                    std::arch::x86_64::_mm_prefetch(
                        ebwt.as_ptr().add(offset + 64) as *const i8,
                        std::arch::x86_64::_MM_HINT_T0,
                    );
                }
            }
        }
    }

    /// Transform this SideLocus to refer to the next side
    pub fn next_side(&mut self, ep: &EbwtParams) {
        assert!(self.valid(), "SideLocus must be valid");
        self.side_byte_off += ep.side_sz as u64;
        self.side_num += 1;
        self.by = 0;
        self.bp = 0;
        self.char_off = 0;
        assert!(self.valid(), "SideLocus must be valid after next_side");
    }

    /// Check if this is a valid (initialized) SideLocus
    pub fn valid(&self) -> bool {
        self.bp != -1
    }

    /// Check if this SideLocus is valid with respect to given parameters
    pub fn is_valid(&self, _ep: &EbwtParams) -> bool {
        self.valid()
    }

    /// Convert locus to BW row it corresponds to
    pub fn to_bw_row(&self) -> u64 {
        self.side_num * Self::SIDE_BWT_LEN + self.char_off as u64
    }

    /// Invalidate this SideLocus (make it look like an invalid one)
    pub fn invalidate(&mut self) {
        self.bp = -1;
    }

    /// Get a read-only pointer to the beginning of the top side
    pub fn side<'a>(&self, ebwt: &'a [u8]) -> &'a [u8] {
        &ebwt[self.side_byte_off as usize..]
    }

    /// Get the actual byte offset in the mmapped data
    pub fn byte_offset(&self) -> u64 {
        self.side_byte_off
    }

    /// Get the mask for extracting the bit-pair at this position
    pub fn mask(&self) -> u8 {
        0x03 << (self.bp * 2)
    }

    /// Extract the nucleotide pair at this locus from data
    pub fn extract(&self, data: &[u8]) -> u8 {
        let byte = data[self.byte_offset() as usize];
        (byte >> (self.bp * 2)) & 0x03
    }
}

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

impl fmt::Display for SideLocus {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "SideLocus(side={}, char_off={}, by={}, bp={})",
            self.side_num, self.char_off, self.by, self.bp
        )
    }
}

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

    #[test]
    fn test_locus_new() {
        let locus = SideLocus::new();
        assert!(!locus.valid());
    }

    #[test]
    fn test_locus_invalid() {
        let mut locus = SideLocus::new();
        assert!(!locus.valid());
        locus.invalidate();
        assert!(!locus.valid());
    }

    #[test]
    fn test_locus_to_bw_row() {
        let locus = SideLocus {
            side_byte_off: 0,
            side_num: 0,
            char_off: 100,
            by: 25,
            bp: 0,
        };
        assert_eq!(locus.to_bw_row(), 100);
    }

    #[test]
    fn test_locus_next_side() {
        let ep = EbwtParams::new(10000);
        let mut locus = SideLocus {
            side_byte_off: ep.side_sz as u64,
            side_num: 1,
            char_off: 50,
            by: 12,
            bp: 2,
        };
        assert!(locus.valid());
        locus.next_side(&ep);
        assert_eq!(locus.side_num, 2);
        assert_eq!(locus.char_off, 0);
        assert_eq!(locus.by, 0);
        assert_eq!(locus.bp, 0);
        assert!(locus.valid());
    }

    #[test]
    fn test_mask() {
        // Test all 4 bit-pair positions
        for bp in 0..4 {
            let locus = SideLocus {
                side_byte_off: 0,
                side_num: 0,
                char_off: 0,
                by: 0,
                bp,
            };
            assert_eq!(locus.mask(), 0x03 << (bp * 2));
        }
    }
}