obikmer/src/obiskio/src/unitig_index/writer.rs

use std::fs::File;
use std::io::{BufWriter, Write as _};
use std::path::Path;

use memmap2::Mmap;
use obikseq::{MAX_KMERS_PER_CHUNK, Unitig};

use crate::error::{SKError, SKResult};

use super::{DEFAULT_BLOCK_BITS, MAGIC, idx_path};

// ── Writer ────────────────────────────────────────────────────────────────────

/// Writes a sequence of [`Unitig`] to an uncompressed binary file and builds
/// a block-sampled offset index at close time.
///
/// One offset is stored every `1 << block_bits` chunks; random access to chunk
/// `i` costs at most `(1 << block_bits) − 1` sequential chunk scans after the
/// block lookup.
///
/// Unitigs with more than [`MAX_KMERS_PER_CHUNK`] k-mers are transparently split
/// into overlapping chunks (k−1 nucleotide overlap) so no k-mer is lost.
pub struct UnitigFileWriter {
    file:          BufWriter<File>,
    block_offsets: Vec<u32>,
    chunk_count:   usize,
    next_offset:   u32,
    n_kmers:       usize,
    k:             usize,
    block_bits:    u8,
    mask:          usize,  // (1 << block_bits) - 1
}

impl UnitigFileWriter {
    /// Create a writer with the default block size (`DEFAULT_BLOCK_BITS = 6`).
    pub fn create(path: &Path) -> SKResult<Self> {
        Self::create_with_block_bits(path, DEFAULT_BLOCK_BITS)
    }

    /// Create a writer with a custom block size.
    ///
    /// `block_bits` must be in 0..=31. `block_bits=0` stores one offset per
    /// chunk (exact, no scan); larger values trade index size for scan length.
    pub fn create_with_block_bits(path: &Path, block_bits: u8) -> SKResult<Self> {
        assert!(block_bits <= 31, "block_bits must be ≤ 31");
        let file = File::create(path).map_err(SKError::Io)?;
        Ok(Self {
            file: BufWriter::new(file),
            block_offsets: Vec::new(),
            chunk_count: 0,
            next_offset: 0,
            n_kmers: 0,
            k: obikseq::params::k(),
            block_bits,
            mask: (1usize << block_bits) - 1,
        })
    }

    /// Write a unitig, splitting into overlapping chunks if it exceeds
    /// [`MAX_KMERS_PER_CHUNK`].
    pub fn write(&mut self, unitig: &Unitig) -> SKResult<()> {
        let seql = unitig.seql();
        let k = self.k;

        if seql < k {
            return Ok(());
        }

        let n_kmers = seql - k + 1;
        if n_kmers <= MAX_KMERS_PER_CHUNK {
            return self.write_chunk(unitig);
        }

        let chunk_nucl = MAX_KMERS_PER_CHUNK + k - 1;
        let stride = MAX_KMERS_PER_CHUNK;
        let mut start = 0;
        while start < seql {
            let end = (start + chunk_nucl).min(seql);
            self.write_chunk(&unitig.sub(start, end))?;
            if end == seql { break; }
            start += stride;
        }
        Ok(())
    }

    fn write_chunk(&mut self, unitig: &Unitig) -> SKResult<()> {
        let seql = unitig.seql();
        let byte_len = (seql + 3) / 4;

        debug_assert!(seql - self.k <= u8::MAX as usize, "chunk exceeds MAX_KMERS_PER_CHUNK");

        if self.chunk_count & self.mask == 0 {
            self.block_offsets.push(self.next_offset);
        }

        self.n_kmers += seql - self.k + 1;
        self.chunk_count += 1;

        unitig.write_to_binary(&mut self.file).map_err(SKError::Io)?;

        self.next_offset += 1 + byte_len as u32;
        Ok(())
    }

    /// Flush and close the binary sequence file.
    ///
    /// The companion `.idx` file is **not** written here; call
    /// [`build_unitig_idx`] separately when exact evidence is needed.
    pub fn close(mut self) -> SKResult<()> {
        self.file.flush().map_err(SKError::Io)?;
        drop(self.file);
        Ok(())
    }

    pub fn len(&self) -> usize { self.chunk_count }
    pub fn is_empty(&self) -> bool { self.chunk_count == 0 }
    pub fn block_bits(&self) -> u8 { self.block_bits }
}

fn write_idx(
    path: &Path,
    n_unitigs: u32,
    n_kmers: u64,
    block_bits: u8,
    block_offsets: &[u32],
) -> SKResult<()> {
    let mut w = BufWriter::new(File::create(path).map_err(SKError::Io)?);
    w.write_all(&MAGIC).map_err(SKError::Io)?;
    w.write_all(&(block_bits as u32).to_le_bytes()).map_err(SKError::Io)?;
    w.write_all(&n_unitigs.to_le_bytes()).map_err(SKError::Io)?;
    w.write_all(&n_kmers.to_le_bytes()).map_err(SKError::Io)?;
    for &off in block_offsets {
        w.write_all(&off.to_le_bytes()).map_err(SKError::Io)?;
    }
    w.flush().map_err(SKError::Io)
}

/// Scan an existing `unitigs.bin` file and write its companion `.idx`.
///
/// Called by the exact-evidence construction route after the sequence file is
/// closed.  `block_bits` controls index granularity (1 << block_bits chunks per
/// offset entry); use [`DEFAULT_BLOCK_BITS`] for the default.
pub fn build_unitig_idx(unitigs_path: &Path, block_bits: u8) -> SKResult<()> {
    assert!(block_bits <= 31, "block_bits must be ≤ 31");

    let file = File::open(unitigs_path).map_err(SKError::Io)?;
    let mmap = unsafe { Mmap::map(&file).map_err(SKError::Io)? };

    let k          = obikseq::params::k();
    let block_size = 1usize << block_bits;
    let mask       = block_size - 1;

    let mut block_offsets: Vec<u32> = Vec::new();
    let mut offset      = 0usize;
    let mut chunk_count = 0usize;
    let mut n_kmers     = 0usize;

    while offset < mmap.len() {
        if chunk_count & mask == 0 {
            block_offsets.push(offset as u32);
        }
        let seql_minus_k = mmap[offset] as usize;
        let byte_len = (seql_minus_k + k + 3) / 4;
        n_kmers     += seql_minus_k + 1;
        offset      += 1 + byte_len;
        chunk_count += 1;
    }

    block_offsets.push(offset as u32); // sentinel

    write_idx(
        &idx_path(unitigs_path),
        chunk_count as u32,
        n_kmers as u64,
        block_bits,
        &block_offsets,
    )
}