slow5lib 0.1.0

use std::cell::RefCell;

use flate2::Compression;
use flate2::read::ZlibDecoder;
use flate2::write::ZlibEncoder;

use crate::Result;
use crate::error::SlowError;
use crate::header::SignalCompression;

/// Decompress a BLOW5 signal field.
///
/// Dispatches on the compression method declared in the file header. For SvbZd,
/// the format is a 4-byte little-endian element count prefix followed by
/// U32Classic-encoded zigzag-delta data -- compatible with the Lemire C
/// streamvbyte library, NOT the ONT VBZ/SVB16 format. For ExZd, the format is
/// svb's self-describing ex-zd frame (no separate length prefix needed).
pub(crate) fn decompress_signal(bytes: &[u8], method: SignalCompression) -> Result<Vec<i16>> {
    match method {
        SignalCompression::None => decompress_signal_none(bytes),
        SignalCompression::SvbZd => decompress_signal_svb_zd(bytes),
        SignalCompression::ExZd => decompress_signal_ex_zd(bytes),
    }
}

/// Maximum decompressed size accepted for a single BLOW5 record.
/// Prevents decompression bombs from hanging fuzz targets or starving real workloads.
pub(crate) const MAX_RECORD_DECOMP: usize = 256 * 1024 * 1024;

/// Decompress a BLOW5 record buffer using a thread-local zstd decompressor context.
///
/// The `ZSTD_DCtx` is allocated once per thread on first use and reused for all
/// subsequent calls on that thread. Safe to call concurrently from multiple threads
/// -- each thread has its own independent context with no shared state.
pub(crate) fn decompress_record_zstd(bytes: &[u8]) -> Result<Vec<u8>> {
    thread_local! {
        static DECOMP: RefCell<Option<zstd::bulk::Decompressor<'static>>> =
            const { RefCell::new(None) };
    }
    DECOMP.with(|cell| {
        let mut opt = cell.borrow_mut();
        let d = opt.get_or_insert_with(|| {
            zstd::bulk::Decompressor::new().expect("zstd decompressor init")
        });
        d.decompress(bytes, MAX_RECORD_DECOMP)
            .map_err(|e| SlowError::Decompression(e.to_string()))
    })
}

/// Decompress a BLOW5 record buffer using zlib.
pub(crate) fn decompress_record_zlib(bytes: &[u8]) -> Result<Vec<u8>> {
    use std::io::Read;
    let mut out = Vec::new();
    ZlibDecoder::new(bytes)
        .take(MAX_RECORD_DECOMP as u64 + 1)
        .read_to_end(&mut out)
        .map_err(|e| SlowError::Decompression(e.to_string()))?;
    if out.len() > MAX_RECORD_DECOMP {
        return Err(SlowError::Decompression(
            "decompressed record exceeds 256 MiB limit".into(),
        ));
    }
    Ok(out)
}

/// Compress a BLOW5 record buffer using a thread-local zstd compressor context.
///
/// Safe to call from multiple rayon workers concurrently -- each thread maintains
/// its own `ZSTD_CCtx` allocated on first use and reused for every subsequent call
/// on that thread.
pub(crate) fn compress_record_zstd_tl(bytes: &[u8]) -> Result<Vec<u8>> {
    thread_local! {
        static COMP: RefCell<Option<zstd::bulk::Compressor<'static>>> =
            const { RefCell::new(None) };
    }
    COMP.with(|cell| {
        let mut opt = cell.borrow_mut();
        let c = opt
            .get_or_insert_with(|| zstd::bulk::Compressor::new(1).expect("zstd compressor init"));
        c.compress(bytes)
            .map_err(|e| SlowError::Decompression(e.to_string()))
    })
}

/// Compress a BLOW5 record buffer using zlib at level 1 (matching C library default).
pub(crate) fn compress_record_zlib(bytes: &[u8]) -> Result<Vec<u8>> {
    use std::io::Write;
    let mut enc = ZlibEncoder::new(Vec::new(), Compression::new(1));
    enc.write_all(bytes)
        .map_err(|e| SlowError::Decompression(e.to_string()))?;
    enc.finish()
        .map_err(|e| SlowError::Decompression(e.to_string()))
}

/// Interpret raw bytes as packed little-endian i16 samples.
fn decompress_signal_none(bytes: &[u8]) -> Result<Vec<i16>> {
    if !bytes.len().is_multiple_of(2) {
        return Err(SlowError::Decompression(format!(
            "uncompressed signal byte count {} is not a multiple of 2",
            bytes.len()
        )));
    }
    Ok(bytes
        .chunks_exact(2)
        .map(|c| i16::from_le_bytes([c[0], c[1]]))
        .collect())
}

/// Decompress slow5 SVB-ZD signal data.
///
/// Wire format: `[n: u32 le][U32Classic(zigzag_delta(widen_i16_to_i32(samples)))]`
///
/// The 4-byte element-count prefix is the caller's responsibility (slow5lib convention);
/// `svb::decode_svbzd_fused` expects the data WITHOUT the prefix and needs `n` separately.
fn decompress_signal_svb_zd(data: &[u8]) -> Result<Vec<i16>> {
    if data.len() < 4 {
        return Err(SlowError::Decompression(
            "SVB-ZD data shorter than 4-byte length prefix".into(),
        ));
    }

    let n = u32::from_le_bytes(data[..4].try_into().unwrap()) as usize;
    // Cap before allocating the output Vec. The SVB decoder allocates n i16 values (2n bytes);
    // an adversarial n field could otherwise cause a multi-GB allocation.
    // 128M samples = 256 MiB of i16 -- far more than any real ONT read.
    const MAX_SIGNAL_SAMPLES: usize = 128 * 1024 * 1024;
    if n > MAX_SIGNAL_SAMPLES {
        return Err(SlowError::Decompression(format!(
            "signal sample count {n} exceeds {MAX_SIGNAL_SAMPLES} limit"
        )));
    }

    svb::decode_svbzd_fused(&data[4..], n).map_err(|e| SlowError::Decompression(e.to_string()))
}

/// Decompress slow5 ex-zd signal data.
///
/// Wire format: svb's self-describing ex-zd frame -- `version(u8) + nin(u64 le) + q(u8)`
/// followed by the patched/exception payload. Unlike SvbZd, no separate length prefix
/// is needed here since the sample count is embedded in the frame itself.
///
/// Uses a thread-local [`svb::ExzdDecoder`] to reuse scratch buffers across calls --
/// BLOW5 files typically hold many thousands of individually-decoded reads, so avoiding
/// a fresh allocation per call matters.
fn decompress_signal_ex_zd(data: &[u8]) -> Result<Vec<i16>> {
    thread_local! {
        static DECODER: RefCell<svb::ExzdDecoder> = RefCell::new(svb::ExzdDecoder::new());
    }
    DECODER.with(|cell| {
        let mut decoder = cell.borrow_mut();
        let mut out = Vec::new();
        decoder
            .decode_into(data, &mut out)
            .map_err(|e| SlowError::Decompression(e.to_string()))?;
        Ok(out)
    })
}

/// Compress signal samples using slow5 ex-zd.
///
/// Wire format: exactly `svb::encode_exzd`'s self-describing frame -- no separate
/// length prefix is added here, unlike SvbZd.
pub(crate) fn compress_signal_ex_zd(samples: &[i16]) -> Result<Vec<u8>> {
    Ok(svb::encode_exzd(samples))
}

/// Compress signal samples using slow5 SVB-ZD.
///
/// Wire format: `[n: u32 le][svb::encode_svbzd(samples)]`
///
/// The 4-byte element-count prefix is prepended here to match the slow5lib wire format.
/// `svb::encode_svbzd` does not include it.
pub(crate) fn compress_signal_svb_zd(samples: &[i16]) -> Result<Vec<u8>> {
    let n = samples.len() as u32;
    let svb_bytes = svb::encode_svbzd(samples);

    let mut out = Vec::with_capacity(4 + svb_bytes.len());
    out.extend_from_slice(&n.to_le_bytes());
    out.extend_from_slice(&svb_bytes);
    Ok(out)
}