gwseq_io/genomic/

chr.rs

//! Chromosome names and sizes, and how a requested name is resolved to one.
//!
//! # The resolution rule
//!
//! [`ChrMap::get`] tries the id, then the `chr`-toggled form of it. Each of
//! those goes through [`ChrMap::try_lookup`], which tries four things in
//! order: the id as given, its lowercase form, its uppercase form, and finally
//! a walk of the map comparing case-insensitively. The walk is there because
//! the three hashed forms only reach an id whose case is uniform — `"ChrX"`
//! misses a file spelling it `"chrX"` — and it takes the first match **in the
//! file's own order**, a file holding both `chrX` and `CHRX` having two, with
//! neither more the answer than the other.
//!
//! # Ids are compared whole
//!
//! Never by prefix, and this is worth stating because the alternative looks
//! reasonable: truncate the query to the length of the longest key first, so
//! that a name longer than the file's fixed-width key field can still match — a
//! bbi chromosome tree stores fixed-width keys, and a writer whose key size is
//! shorter than a name stores that name truncated. But the key size cannot be
//! inferred from the map. The longest *present* name is only a lower bound on
//! the file's declared key size, so a file holding `chr1` alone reads as having
//! a key size of 4, and every longer query is then answered with whichever
//! chromosome it shares a prefix with — `12` resolving to `chr1`, yeast's
//! `chrXVII` to `chrXVI`. Writes take the same path, so writing to `12` against
//! `chr_sizes={"chr1": ...}` would land on chr1.
//!
//! The declared key size is still threaded through, for the *error message*
//! only: a lookup that fails against a file whose key field is narrower than
//! the query says so, and names the chromosome the query's first `key_size`
//! characters spell when the file holds one.

use indexmap::IndexMap;

use crate::error::{Error, Result};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ChrEntry {
    pub id: String,
    pub size: i64,
    /// 0-based position in the file's reference order, as bbi and BAM store it.
    pub index: usize,
}

/// Insertion-ordered chromosome map. Order is the file's reference order, which
/// callers depend on — `chr_sizes` is documented as coming back in it, and the
/// case-insensitive walk in `try_lookup` — private, so not linked — resolves
/// ties by it.
#[derive(Debug, Clone, Default)]
pub struct ChrMap {
    entries: IndexMap<String, ChrEntry>,
    /// The key width the file declares, or `None` when there is no file behind
    /// this map. Used in error messages only — never in matching.
    declared_key_size: Option<usize>,
}

impl ChrMap {
    /// Build from names and sizes, numbering them by insertion order.
    ///
    /// A name given twice keeps the index it was first given, and its later
    /// size wins. Numbering from `map.len()` instead left the two entries
    /// sharing an index — the second overwrote the first in the map while the
    /// count stood still — and an index is what a data record carries.
    pub fn from_entries(entries: impl IntoIterator<Item = (String, i64)>) -> Self {
        let mut map: IndexMap<String, ChrEntry> = IndexMap::new();
        for (id, size) in entries {
            let index = map.get(&id).map_or(map.len(), |e| e.index);
            map.insert(id.clone(), ChrEntry { id, size, index });
        }
        Self {
            entries: map,
            declared_key_size: None,
        }
    }

    /// Build from names, sizes and the indices **the file gives them**.
    ///
    /// A bbi chromosome tree and a BAM reference list both store an explicit
    /// index, and that index is what data records and R-tree items carry — so
    /// it cannot be re-derived from position here. Entries are sorted by it,
    /// which is the order `chr_sizes` reports them in.
    pub fn from_indexed_entries(entries: impl IntoIterator<Item = (String, i64, usize)>) -> Self {
        let mut list: Vec<ChrEntry> = entries
            .into_iter()
            .map(|(id, size, index)| ChrEntry { id, size, index })
            .collect();
        list.sort_by_key(|e| e.index);
        let mut map = IndexMap::with_capacity(list.len());
        for entry in list {
            map.insert(entry.id.clone(), entry);
        }
        Self {
            entries: map,
            declared_key_size: None,
        }
    }

    /// The fixed key width a bbi chromosome tree declares. Message-only.
    pub fn with_key_size(mut self, key_size: usize) -> Self {
        self.declared_key_size = Some(key_size);
        self
    }

    pub fn declared_key_size(&self) -> Option<usize> {
        self.declared_key_size
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    pub fn iter(&self) -> impl Iterator<Item = &ChrEntry> {
        self.entries.values()
    }

    pub fn names(&self) -> Vec<String> {
        self.entries.keys().cloned().collect()
    }

    /// The chromosome the **file** numbers `index`, not the one at that
    /// position.
    ///
    /// The two coincide on every file whose indices are dense and 0-based,
    /// which is every file any writer here or at UCSC produces — so the
    /// positional lookup is tried first and the walk is the fallback for a file
    /// that numbers its chromosomes otherwise.
    pub fn by_index(&self, index: usize) -> Option<&ChrEntry> {
        if let Some((_, entry)) = self.entries.get_index(index) {
            if entry.index == index {
                return Some(entry);
            }
        }
        self.entries.values().find(|e| e.index == index)
    }

    /// Total size of every chromosome, which is what a whole-genome walk covers.
    pub fn genome_size(&self) -> i64 {
        self.entries.values().map(|e| e.size).sum()
    }

    /// The id as given, then its lowercase form, then its uppercase form, then
    /// a case-insensitive walk in file order.
    fn try_lookup(&self, id: &str) -> Option<&ChrEntry> {
        if let Some(entry) = self.entries.get(id) {
            return Some(entry);
        }
        if let Some(entry) = self.entries.get(&id.to_ascii_lowercase()) {
            return Some(entry);
        }
        if let Some(entry) = self.entries.get(&id.to_ascii_uppercase()) {
            return Some(entry);
        }
        self.entries
            .iter()
            .find(|(key, _)| key.len() == id.len() && key.eq_ignore_ascii_case(id))
            .map(|(_, entry)| entry)
    }

    /// The `chr`-toggled form: `chr1` ↔ `1`.
    ///
    /// `get(..3)` rather than `&id[..3]`, which would panic on an id whose
    /// third byte is inside a multi-byte character.
    fn alt_id(id: &str) -> String {
        match id.get(..3) {
            Some(prefix) if prefix.eq_ignore_ascii_case("chr") => id[3..].to_string(),
            _ => format!("chr{id}"),
        }
    }

    pub fn get(&self, id: &str) -> Option<&ChrEntry> {
        if let Some(entry) = self.try_lookup(id) {
            return Some(entry);
        }
        self.try_lookup(&Self::alt_id(id))
    }

    pub fn contains(&self, id: &str) -> bool {
        self.get(id).is_some()
    }

    /// Resolve a requested name, or fail with everything known about why not.
    pub fn resolve(&self, id: &str) -> Result<&ChrEntry> {
        match self.get(id) {
            Some(entry) => Ok(entry),
            None => Err(self.not_found(id)),
        }
    }

    /// The error a failed lookup carries, as parts rather than a sentence.
    ///
    /// `key_size` is set only when the id cannot fit the file's declared name
    /// field — which is what makes the lookup impossible rather than merely
    /// unsuccessful — and `truncated_match` only when the id's first
    /// `key_size` characters do spell a chromosome the file holds. Both forms
    /// of the id are tried, since "chrX" may not fit a field that "X" would.
    /// [`crate::error::Error`] turns them into the sentence.
    fn not_found(&self, id: &str) -> Error {
        let alt = Self::alt_id(id);
        let mut key_size = None;
        let mut truncated_match = None;
        if let Some(declared) = self.declared_key_size.filter(|&k| k > 0) {
            if id.chars().count() > declared && alt.chars().count() > declared {
                key_size = Some(declared);
                truncated_match = [id, alt.as_str()]
                    .into_iter()
                    .filter_map(|candidate| candidate.get(..declared))
                    .find_map(|prefix| self.try_lookup(prefix))
                    .map(|entry| entry.id.clone());
            }
        }
        Error::UnknownChromosome {
            id: id.to_string(),
            key_size,
            truncated_match,
            available: self.entries.keys().cloned().collect(),
        }
    }

    /// Which chromosomes a request walks: the ones named, resolved, in the
    /// order they were asked for — or every chromosome in the file's order when
    /// none were named.
    pub fn select(&self, requested: &[String]) -> Result<Vec<ChrEntry>> {
        if requested.is_empty() {
            return Ok(self.entries.values().cloned().collect());
        }
        requested
            .iter()
            .map(|id| self.resolve(id).cloned())
            .collect()
    }
}

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

    fn map(names: &[(&str, i64)]) -> ChrMap {
        ChrMap::from_entries(names.iter().map(|(n, s)| (n.to_string(), *s)))
    }

    #[test]
    fn resolves_exact_case_and_prefix_toggle() {
        let m = map(&[("chr1", 100), ("chr2", 200)]);
        for asked in ["chr1", "CHR1", "Chr1", "1"] {
            assert_eq!(m.get(asked).unwrap().id, "chr1", "asking for {asked}");
        }
        let m = map(&[("1", 100), ("2", 200)]);
        for asked in ["1", "chr1", "CHR1"] {
            assert_eq!(m.get(asked).unwrap().id, "1", "asking for {asked}");
        }
    }

    #[test]
    fn mixed_case_id_reaches_a_uniform_case_key() {
        // "ChrX" is neither all-lower nor all-upper, so only the walk finds it.
        let m = map(&[("chrX", 10)]);
        assert_eq!(m.get("ChrX").unwrap().id, "chrX");
        assert_eq!(m.get("X").unwrap().id, "chrX");
    }

    #[test]
    fn the_uppercase_form_is_tried_before_the_walk() {
        // "ChrX" uppercases to "CHRX", which this map has, so the hashed
        // lookup answers and the walk never runs.
        let m = map(&[("chrX", 1), ("CHRX", 2)]);
        assert_eq!(m.get("ChrX").unwrap().id, "CHRX");
    }

    #[test]
    fn ties_in_the_walk_go_to_file_order() {
        // Neither "chrx" nor "CHRX" is present, so only the walk can match,
        // and it takes the first of the two in the file's own order.
        let m = map(&[("chrX", 1), ("cHRx", 2)]);
        assert_eq!(m.get("ChrX").unwrap().id, "chrX");
        let m = map(&[("cHRx", 1), ("chrX", 2)]);
        assert_eq!(m.get("ChrX").unwrap().id, "cHRx");
    }

    #[test]
    fn ids_are_compared_whole() {
        // Prefix matching would resolve "12" to "chr1", "chrXVII" to "chrXVI".
        let m = map(&[("chr1", 100)]);
        assert!(m.get("chr12").is_none());
        assert!(m.get("12").is_none());
        let m = map(&[("chrXVI", 100)]);
        assert!(m.get("chrXVII").is_none());
    }

    #[test]
    fn not_found_message_lists_what_is_available() {
        let m = map(&[("chr1", 1), ("chr2", 2)]);
        let message = m.resolve("chrZ").unwrap_err().to_string();
        assert!(message.contains("not found"), "{message}");
        assert!(message.contains("(available: chr1, chr2)"), "{message}");
        assert!(!message.contains("field"), "{message}");
    }

    #[test]
    fn an_over_long_id_is_told_it_cannot_fit() {
        let m = map(&[("chr1_GL456210_random", 1)]).with_key_size(20);
        // 24 characters, and stripping "chr" still leaves 21 — over the field
        // both ways, which is what makes the clause apply.
        let message = m
            .resolve("chrX_GL456210_random_ext")
            .unwrap_err()
            .to_string();
        assert!(message.contains("20-character field"), "{message}");
        assert!(message.contains("24 characters long"), "{message}");
    }

    #[test]
    fn an_over_long_id_names_the_chromosome_its_prefix_spells() {
        let m = map(&[("chr1_GL456210_random", 1)]).with_key_size(20);
        // 25 characters, 22 once "chr" comes off: over the field both ways, and
        // its first 20 characters spell a chromosome the file does hold.
        let message = m
            .resolve("chr1_GL456210_random_v234")
            .unwrap_err()
            .to_string();
        assert!(
            message.contains("does hold chr1_GL456210_random"),
            "{message}"
        );
    }

    #[test]
    fn an_id_that_fits_once_stripped_gets_no_field_clause() {
        // 23 characters, but "chr" comes off and 20 fits the field, so the
        // file could have held it and the clause must not fire.
        let m = map(&[("chr1_GL456210_random", 1)]).with_key_size(20);
        let message = m
            .resolve("chrX_GL456210_random_ab")
            .unwrap_err()
            .to_string();
        assert!(!message.contains("field"), "{message}");
        assert!(message.contains("not found"), "{message}");
    }

    #[test]
    fn select_takes_request_order_and_defaults_to_file_order() {
        let m = map(&[("chr1", 1), ("chr2", 2), ("chr3", 3)]);
        let all: Vec<_> = m
            .select(&[])
            .unwrap()
            .iter()
            .map(|e| e.id.clone())
            .collect();
        assert_eq!(all, ["chr1", "chr2", "chr3"]);
        let some: Vec<_> = m
            .select(&["chr3".into(), "1".into()])
            .unwrap()
            .iter()
            .map(|e| e.id.clone())
            .collect();
        assert_eq!(some, ["chr3", "chr1"]);
    }

    #[test]
    fn index_follows_insertion_order() {
        let m = map(&[("chr1", 1), ("chr2", 2)]);
        assert_eq!(m.get("chr2").unwrap().index, 1);
        assert_eq!(m.by_index(0).unwrap().id, "chr1");
        assert_eq!(m.genome_size(), 3);
    }

    #[test]
    fn file_indices_are_kept_and_sorted_by() {
        // Deliberately out of order and not 0-based, which is what a positional
        // by_index would get wrong.
        let m = ChrMap::from_indexed_entries([
            ("chrB".to_string(), 20, 7),
            ("chrA".to_string(), 10, 3),
        ]);
        assert_eq!(m.names(), ["chrA", "chrB"]);
        assert_eq!(m.get("chrA").unwrap().index, 3);
        assert_eq!(m.by_index(3).unwrap().id, "chrA");
        assert_eq!(m.by_index(7).unwrap().id, "chrB");
        assert!(m.by_index(0).is_none());
    }
}