rustalign_fmindex/

reference.rs

//! Reference sequence storage and extraction
//!
//! This module handles loading the reference sequence from .3.rai and .4.rai
//! files and extracting stretches for Smith-Waterman alignment.

use memmap2::Mmap;
use rustalign_common::{AlignError, AlignResult, Nuc};
use std::fs::File;
use std::path::Path;

/// A record describing an unambiguous stretch of the reference
#[derive(Debug, Clone)]
struct RefRecord {
    /// Number of ambiguous (N) characters before this stretch
    off: u32,
    /// Length of the unambiguous stretch
    len: u32,
    /// Whether this is the first record for a reference sequence
    #[allow(dead_code)]
    first: bool,
}

/// Reference sequence storage
///
/// Loads and stores reference sequences from .3.rai and .4.rai RustAlign index files.
/// The reference is stored in bit-packed format (2 bits per nucleotide).
pub struct Reference {
    /// Records describing unambiguous stretches
    recs: Vec<RefRecord>,

    /// Cumulative count of unambiguous characters up to each record
    cum_unambig: Vec<u64>,

    /// Approximate lengths of reference sequences (total chars including Ns)
    ref_lens: Vec<u64>,

    /// Buffer begin offsets per reference sequence (unambiguous chars only)
    ref_offs: Vec<u64>,

    /// Record begin/end indices per reference sequence
    ref_rec_offs: Vec<(usize, usize)>,

    /// Number of reference sequences
    n_refs: u32,

    /// Total size of buffer needed
    buf_sz: u64,
}

impl Reference {
    /// Load reference from .3.rai and .4.rai files
    ///
    /// # Arguments
    /// * `path_base` - Path to the index base (without extensions)
    pub fn load<P: AsRef<Path>>(path_base: P) -> AlignResult<Self> {
        let path_base = path_base.as_ref();
        let base_name = path_base
            .file_name()
            .ok_or_else(|| AlignError::InvalidFormat("Invalid index path".into()))?
            .to_str()
            .unwrap();

        let file_3 = path_base.with_file_name(format!("{}.3.rai", base_name));

        if !file_3.exists() {
            return Err(AlignError::IndexCorrupted {
                path: path_base.to_path_buf(),
            });
        }

        // Load .3.rai (metadata - RefRecords)
        let file3 = File::open(&file_3)?;
        let mmap3 = unsafe { Mmap::map(&file3)? };

        Self::parse_from_mmap(&mmap3)
    }

    /// Parse reference from memory-mapped .3.rai file
    fn parse_from_mmap(meta: &[u8]) -> AlignResult<Self> {
        use byteorder::{LittleEndian, ReadBytesExt};
        use std::io::Cursor;

        if meta.len() < 8 {
            return Err(AlignError::InvalidFormat(
                "Reference metadata too small".into(),
            ));
        }

        let mut cursor = Cursor::new(meta);

        // Read endian hint
        let endian_hint = cursor
            .read_i32::<LittleEndian>()
            .map_err(|e| AlignError::InvalidFormat(format!("Failed to read endian hint: {}", e)))?;

        if endian_hint != 1 {
            return Err(AlignError::InvalidFormat(format!(
                "Invalid endian hint: {} (expected 1)",
                endian_hint
            )));
        }

        // Read number of records (sz)
        let num_records = cursor
            .read_u32::<LittleEndian>()
            .map_err(|e| AlignError::InvalidFormat(format!("Failed to read record count: {}", e)))?
            as usize;

        if num_records == 0 {
            return Err(AlignError::InvalidFormat(
                "Number of reference records is 0".into(),
            ));
        }

        // Parse RefRecords
        // Each record is: off (u32), len (u32), first (u8)
        let mut recs = Vec::with_capacity(num_records);
        let mut cum_unambig = Vec::with_capacity(num_records);
        let mut ref_lens = Vec::new();
        let mut ref_offs = Vec::new();
        let mut ref_rec_offs = Vec::new();

        let mut n_refs = 0u32;
        let mut cumsz: u64 = 0; // cumulative unambiguous chars
        let mut cumlen: u64 = 0; // cumulative total chars for current ref

        for _ in 0..num_records {
            let pos = cursor.position() as usize;
            if pos + 9 > meta.len() {
                break;
            }

            // Read RefRecord: off (u32), len (u32), first (u8)
            let off = cursor.read_u32::<LittleEndian>().map_err(|e| {
                AlignError::InvalidFormat(format!("Failed to read record off: {}", e))
            })?;

            let len = cursor.read_u32::<LittleEndian>().map_err(|e| {
                AlignError::InvalidFormat(format!("Failed to read record len: {}", e))
            })?;

            let first_byte = cursor.read_u8().map_err(|e| {
                AlignError::InvalidFormat(format!("Failed to read record first: {}", e))
            })?;
            let first = first_byte != 0;

            if first {
                // This is the first record for a new reference sequence
                ref_rec_offs.push((recs.len(), recs.len()));
                ref_offs.push(cumsz);
                if n_refs > 0 {
                    // Close out the previous reference
                    ref_lens.push(cumlen);
                    ref_rec_offs[n_refs as usize - 1].1 = recs.len();
                }
                cumlen = 0;
                n_refs += 1;
            } else if recs.is_empty() {
                return Err(AlignError::InvalidFormat(
                    "First record in reference index was not marked as 'first'".into(),
                ));
            }

            cum_unambig.push(cumsz);
            cumsz += len as u64;
            cumlen += off as u64; // count ambiguous chars before
            cumlen += len as u64; // count unambiguous chars

            recs.push(RefRecord { off, len, first });
        }

        // Store cap entries for the last reference
        if n_refs > 0 {
            ref_rec_offs[n_refs as usize - 1].1 = recs.len();
        }
        ref_rec_offs.push((recs.len(), recs.len()));
        ref_offs.push(cumsz);
        ref_lens.push(cumlen);
        cum_unambig.push(cumsz);

        Ok(Self {
            recs,
            cum_unambig,
            ref_lens,
            ref_offs,
            ref_rec_offs,
            n_refs,
            buf_sz: cumsz,
        })
    }

    /// Get the number of reference sequences
    pub fn n_refs(&self) -> u32 {
        self.n_refs
    }

    /// Get the approximate length of a reference sequence
    pub fn approx_len(&self, tidx: u32) -> u64 {
        if (tidx as usize) < self.ref_lens.len() {
            self.ref_lens[tidx as usize]
        } else {
            0
        }
    }

    /// Extract a stretch of reference sequence
    ///
    /// # Arguments
    /// * `buf` - The bit-packed reference buffer from .4.rai file
    /// * `tidx` - Reference sequence index (chromosome index)
    /// * `toff` - Offset within the reference sequence
    /// * `count` - Number of bases to extract
    ///
    /// # Returns
    /// A vector of nucleotides
    pub fn get_stretch(
        &self,
        buf: &[u8],
        tidx: u32,
        toff: u64,
        count: usize,
    ) -> AlignResult<Vec<Nuc>> {
        if (tidx as usize) >= self.ref_rec_offs.len() {
            return Err(AlignError::InvalidArgument(format!(
                "Invalid reference index: {}",
                tidx
            )));
        }

        let (rec_begin, rec_end) = self.ref_rec_offs[tidx as usize];

        // Walk through records to find the right position
        let mut cur_off: u64 = 0; // Current offset within this reference
        let mut buf_off: u64 = self.ref_offs.get(tidx as usize).copied().unwrap_or(0);

        for reci in rec_begin..rec_end {
            let rec = &self.recs[reci];
            cur_off += rec.off as u64; // Skip ambiguous chars

            // Check if toff falls within this record's unambiguous stretch
            if toff < cur_off + rec.len as u64 {
                // Found the right record
                if toff >= cur_off {
                    // toff is within the unambiguous part
                    let within_rec = toff - cur_off;
                    buf_off += within_rec;

                    // Calculate byte/bit offset in bit-packed buffer
                    let byte_off = (buf_off / 4) as usize;
                    let bit_off = (buf_off % 4) * 2;

                    // Extract nucleotides
                    let mut result = Vec::with_capacity(count);
                    let mut current_byte_off = byte_off;
                    let mut current_bit_off = bit_off as u8;
                    let mut remaining_in_rec = (rec.len as u64 - within_rec) as usize;
                    let mut next_rec = reci + 1;

                    for _ in 0..count {
                        if remaining_in_rec == 0 {
                            // Need to move to next record
                            if next_rec < rec_end {
                                // Skip ambiguous chars (treat as N)
                                let next_rec_obj = &self.recs[next_rec];
                                if next_rec_obj.off > 0 {
                                    result.push(Nuc::N);
                                    // For simplicity, just continue with Ns for ambiguous regions
                                    // In a full implementation, we'd handle this more carefully
                                }
                                remaining_in_rec = next_rec_obj.len as usize;
                                current_byte_off = (self.cum_unambig[next_rec] / 4) as usize;
                                current_bit_off = (self.cum_unambig[next_rec] % 4 * 2) as u8;
                                next_rec += 1;
                            } else {
                                // End of reference
                                result.push(Nuc::N);
                                continue;
                            }
                        }

                        if current_byte_off >= buf.len() {
                            result.push(Nuc::N);
                            continue;
                        }

                        let byte = buf[current_byte_off];
                        let nuc_code = (byte >> current_bit_off) & 0x03;

                        let nuc = match nuc_code {
                            0 => Nuc::A,
                            1 => Nuc::C,
                            2 => Nuc::G,
                            3 => Nuc::T,
                            _ => Nuc::N,
                        };
                        result.push(nuc);

                        // Advance to next nucleotide
                        current_bit_off += 2;
                        if current_bit_off >= 8 {
                            current_bit_off = 0;
                            current_byte_off += 1;
                        }
                        remaining_in_rec -= 1;
                    }

                    return Ok(result);
                }
            }

            cur_off += rec.len as u64;
            buf_off += rec.len as u64;
        }

        // toff is in an ambiguous region at the end
        Ok(vec![Nuc::N; count])
    }

    /// Get the buffer size needed
    pub fn buf_sz(&self) -> u64 {
        self.buf_sz
    }
}

/// Combined structure for FM-index with reference
#[allow(dead_code)]
pub struct EbwtWithRef {
    /// The FM-index
    pub ebwt: super::Ebwt,
    /// The reference sequence metadata
    pub reference: Reference,
    /// The memory-mapped reference buffer
    ref_buf: Mmap,
    /// Mapping from FM-index chromosome index to Reference chromosome index
    /// FM-index uses rstarts array ordering, Reference uses file ordering
    fm_to_ref_chr_map: Vec<u32>,
}

impl EbwtWithRef {
    /// Load both FM-index and reference
    pub fn load<P: AsRef<Path>>(path_base: P) -> AlignResult<Self> {
        let path_base = path_base.as_ref();
        let base_name = path_base
            .file_name()
            .ok_or_else(|| AlignError::InvalidFormat("Invalid index path".into()))?
            .to_str()
            .unwrap();

        // Load the FM-index
        let ebwt = super::Ebwt::load(path_base)?;

        // Load the reference metadata
        let reference = Reference::load(path_base)?;

        // Memory-map the .4.rai file for the reference buffer
        let file_4 = path_base.with_file_name(format!("{}.4.rai", base_name));
        let file = File::open(&file_4)?;
        let ref_buf = unsafe { Mmap::map(&file)? };

        Ok(Self {
            ebwt,
            reference,
            ref_buf,
            fm_to_ref_chr_map: Vec::new(), // Will build below
        })
    }

    /// Extract reference sequence for a given chromosome and position
    ///
    /// Note: pos should be 0-based position within the chromosome
    pub fn get_reference(&self, chr_idx: u32, pos: u64, len: usize) -> AlignResult<Vec<Nuc>> {
        self.reference.get_stretch(&self.ref_buf, chr_idx, pos, len)
    }

    /// Get the underlying Ebwt
    pub fn ebwt(&self) -> &super::Ebwt {
        &self.ebwt
    }
}