SimurgArray/rust/simurg-core/src/bpm.rs

//! BPM LocusEntry records — port of `simurg.formats.bpm._parse_locus_entries` (format: Illumina BeadArrayFiles
//! `BeadPoolManifest.py`, BSD-2). Same field order, checks and error messages; strings of ≤ 64 bytes are shared
//! Python objects, as in the Python parser.

use std::collections::HashMap;

use numpy::{IntoPyArray, PyArray1, PyArray2};
use pyo3::exceptions::PyValueError;
use pyo3::prelude::*;
use pyo3::types::{PyList, PyString};

pub const STR_COLS: usize = 16;
// output column order (matches the Python `str_cols` list)
const ILMN: usize = 0;
const SNP: usize = 1;
const CHR: usize = 2;
const ISTR: usize = 3;
const SSTR: usize = 4;
const RSTR: usize = 5;
const PLOIDY: usize = 6;
const SPECIES: usize = 7;
const BUILD: usize = 8;
const SOURCE: usize = 9;
const SVER: usize = 10;
const CSTR: usize = 11;
const SSEQ: usize = 12;
const TOP: usize = 13;
const PA: usize = 14;
const PB: usize = 15;

struct Cursor<'a> {
    b: &'a [u8],
    pos: usize,
}

fn err(msg: impl Into<String>) -> PyErr {
    PyValueError::new_err(msg.into())
}

impl<'a> Cursor<'a> {
    fn string(&mut self) -> PyResult<&'a [u8]> {
        let size = self.b.len();
        let mut p = self.pos;
        if p >= size {
            return Err(err("unexpected end of file in locus entry"));
        }
        let mut byte = self.b[p];
        p += 1;
        let mut ln = (byte & 0x7f) as usize;
        let mut shift = 7;
        while byte & 0x80 != 0 {
            if p >= size || shift > 28 {
                return Err(err("bad string length in locus entry"));
            }
            byte = self.b[p];
            p += 1;
            ln |= ((byte & 0x7f) as usize) << shift;
            shift += 7;
        }
        let end = p + ln;
        if end > size {
            return Err(err("string runs past end of file in locus entry"));
        }
        self.pos = end;
        Ok(&self.b[p..end])
    }
    fn i32(&mut self, what: &str) -> PyResult<i32> {
        if self.pos + 4 > self.b.len() {
            return Err(err(format!("unexpected end of file at {what}")));
        }
        let v = i32::from_le_bytes(self.b[self.pos..self.pos + 4].try_into().unwrap());
        self.pos += 4;
        Ok(v)
    }
}

type Parsed<'py> = (
    Vec<Bound<'py, PyList>>,
    Bound<'py, PyArray1<i64>>,
    Bound<'py, PyArray1<i32>>,
    Bound<'py, PyArray1<i32>>,
    Bound<'py, PyArray1<i16>>,
    Bound<'py, PyArray1<i8>>,
    Bound<'py, PyArray1<bool>>,
    Bound<'py, PyArray1<i8>>,
    Option<Bound<'py, PyArray2<f32>>>,
    usize,
);

/// Parse `n` LocusEntry records starting at byte 0 of `buf`. Returns the 16 string columns, map_infos,
/// address_a, address_b, assay types (−1 = infer from sequences, entry version < 6), expected clusters,
/// intensity-only flags, entry versions, base fractions (None unless some entry is v7+) and the bytes consumed.
/// Errors are ValueError with the Python parser's message.
#[pyfunction]
pub fn bpm_locus_entries<'py>(py: Python<'py>, buf: &[u8], n: usize, names: Vec<String>) -> PyResult<Parsed<'py>> {
    let name_to_idx: HashMap<&str, usize> = names.iter().enumerate().map(|(i, s)| (s.as_str(), i)).collect();
    let mut cols: Vec<Vec<Option<Py<PyString>>>> = (0..STR_COLS).map(|_| (0..n).map(|_| None).collect()).collect();
    let mut intern: HashMap<&[u8], Py<PyString>> = HashMap::new();
    let mut map_infos = vec![0i64; n];
    let mut address_a = vec![0i32; n];
    let mut address_b = vec![0i32; n];
    let mut assay = vec![0i16; n];
    let mut exp_clusters = vec![0i8; n];
    let mut int_only = vec![false; n];
    let mut versions = vec![0i8; n];
    let mut fracs = vec![0f32; 4 * n];
    let mut seen = vec![false; n];
    let mut any_v7 = false;
    let mut c = Cursor { b: buf, pos: 0 };

    let to_py = |raw: &'_ [u8], intern: &mut HashMap<&[u8], Py<PyString>>| -> Py<PyString> {
        if let Some(s) = intern.get(raw) {
            return s.clone_ref(py);
        }
        let s = PyString::new(py, &String::from_utf8_lossy(raw)).unbind();
        s
    };

    for _ in 0..n {
        let ver = c.i32("locus entry")?;
        if ![4, 6, 7, 8].contains(&ver) {
            return Err(err(format!("unsupported locus entry version {ver}")));
        }
        let ilmn_id = c.string()?;
        let name = c.string()?;
        for _ in 0..3 {
            c.string()?;
        }
        c.i32("locus index")?;
        c.string()?;
        let ilmn_strand = c.string()?;
        let snp = c.string()?;
        let chrom = c.string()?;
        let ploidy = c.string()?;
        let species = c.string()?;
        let map_info = c.string()?;
        let top_seq = c.string()?;
        let cust_strand = c.string()?;
        if c.pos + 8 > buf.len() {
            return Err(err("unexpected end of file at addresses"));
        }
        let a_addr = c.i32("addresses")?;
        let b_addr = c.i32("addresses")?;
        let probe_a = c.string()?;
        let probe_b = c.string()?;
        let build = c.string()?;
        let source = c.string()?;
        let source_ver = c.string()?;
        let source_strand = c.string()?;
        let source_seq = c.string()?;
        let (mut at, mut exp_cl, mut io) = (-1i16, 0i8, 0u8);
        if ver >= 6 {
            if c.pos + 4 > buf.len() {
                return Err(err("unexpected end of file at assay type"));
            }
            exp_cl = buf[c.pos + 1] as i8;
            io = buf[c.pos + 2];
            at = buf[c.pos + 3] as i16;
            c.pos += 4;
            if at > 2 {
                return Err(err(format!("invalid assay type {at}")));
            }
            if (b_addr == 0) != (at == 0) {
                return Err(err("assay type is inconsistent with address B"));
            }
        }
        let mut f4: Option<[f32; 4]> = None;
        if ver >= 7 {
            if c.pos + 16 > buf.len() {
                return Err(err("unexpected end of file at base fractions"));
            }
            let mut f = [0f32; 4];
            for (k, v) in f.iter_mut().enumerate() {
                *v = f32::from_le_bytes(buf[c.pos + 4 * k..c.pos + 4 * k + 4].try_into().unwrap());
            }
            f4 = Some(f);
            c.pos += 16;
        }
        let ref_strand: &[u8] = if ver >= 8 { c.string()? } else { b"" };

        let name_s = String::from_utf8_lossy(name);
        let idx = match name_to_idx.get(name_s.as_ref()) {
            Some(&i) => i,
            None => return Err(err(format!("locus entry {} not in the names table", py_repr(&name_s)))),
        };
        if seen[idx] {
            return Err(err(format!("locus entry {} appears twice", py_repr(&name_s))));
        }
        seen[idx] = true;
        for (col, raw) in [
            (ILMN, ilmn_id),
            (SNP, snp),
            (CHR, chrom),
            (ISTR, ilmn_strand),
            (SSTR, source_strand),
            (RSTR, ref_strand),
            (PLOIDY, ploidy),
            (SPECIES, species),
            (BUILD, build),
            (SOURCE, source),
            (SVER, source_ver),
            (CSTR, cust_strand),
            (SSEQ, source_seq),
            (TOP, top_seq),
            (PA, probe_a),
            (PB, probe_b),
        ] {
            let s = to_py(raw, &mut intern);
            if raw.len() <= 64 && !intern.contains_key(raw) {
                intern.insert(raw, s.clone_ref(py));
            }
            cols[col][idx] = Some(s);
        }
        map_infos[idx] = if map_info.is_empty() {
            0
        } else {
            let t = String::from_utf8_lossy(map_info);
            match t.trim().parse::<i64>() {
                Ok(v) => v,
                Err(_) => {
                    return Err(err(format!(
                        "non-integer MapInfo {} for {}",
                        py_repr(&t),
                        py_repr(&name_s)
                    )))
                }
            }
        };
        address_a[idx] = a_addr;
        address_b[idx] = b_addr;
        assay[idx] = at;
        exp_clusters[idx] = exp_cl;
        int_only[idx] = io != 0;
        versions[idx] = ver as i8;
        if let Some(f) = f4 {
            fracs[4 * idx..4 * idx + 4].copy_from_slice(&f);
            any_v7 = true;
        }
    }
    let missing = seen.iter().filter(|s| !**s).count();
    if missing > 0 {
        return Err(err(format!("{missing} loci in the names table have no locus entry")));
    }
    let empty = PyString::new(py, "").unbind();
    let lists = cols
        .into_iter()
        .map(|col| PyList::new(py, col.into_iter().map(|s| s.unwrap_or_else(|| empty.clone_ref(py)))))
        .collect::<PyResult<Vec<_>>>()?;
    let fr = if any_v7 {
        Some(numpy::ndarray::Array2::from_shape_vec((n, 4), fracs).map_err(|e| err(e.to_string()))?.into_pyarray(py))
    } else {
        None
    };
    Ok((
        lists,
        map_infos.into_pyarray(py),
        address_a.into_pyarray(py),
        address_b.into_pyarray(py),
        assay.into_pyarray(py),
        exp_clusters.into_pyarray(py),
        int_only.into_pyarray(py),
        versions.into_pyarray(py),
        fr,
        c.pos,
    ))
}

/// Python `repr()` of a str for error messages (single quotes unless the text contains one).
fn py_repr(s: &str) -> String {
    if s.contains('\'') && !s.contains('"') {
        format!("\"{s}\"")
    } else {
        format!("'{}'", s.replace('\\', "\\\\").replace('\'', "\\'"))
    }
}