molframe_query/
residue_ranges.rs

1//! Residue numbers written as ranges, and as a `resid` query over them.
2
3use std::collections::BTreeMap;
4
5/// Consecutive runs `(first, last)` of the sorted, distinct `numbers`.
6fn runs(numbers: &[i32]) -> Vec<(i32, i32)> {
7    let mut out: Vec<(i32, i32)> = Vec::new();
8    for &number in numbers {
9        match out.last_mut() {
10            Some((_, end)) if end.checked_add(1) == Some(number) => *end = number,
11            _ => out.push((number, number)),
12        }
13    }
14    out
15}
16
17fn sorted_distinct(numbers: &[i32]) -> Vec<i32> {
18    let mut sorted = numbers.to_vec();
19    sorted.sort_unstable();
20    sorted.dedup();
21    sorted
22}
23
24/// `[20, 21, 22, 25]` as `20-22,25`. The input need not be sorted and may repeat.
25#[must_use]
26pub fn compact_ranges(numbers: &[i32]) -> String {
27    runs(&sorted_distinct(numbers))
28        .into_iter()
29        .map(|(start, end)| {
30            if start == end {
31                start.to_string()
32            } else {
33                format!("{start}-{end}")
34            }
35        })
36        .collect::<Vec<_>>()
37        .join(",")
38}
39
40/// A `resid` query for residues given as (chain, number) pairs: consecutive numbers share a
41/// range (`resid 3:7`), and each chain is its own scope (an empty chain name has none). Returns
42/// the query and the number of distinct residues, or `None` when there are none.
43#[must_use]
44pub fn resid_query(residues: &[(String, i32)]) -> Option<(String, usize)> {
45    let mut by_chain: BTreeMap<&str, Vec<i32>> = BTreeMap::new();
46    for (chain, number) in residues {
47        by_chain.entry(chain).or_default().push(*number);
48    }
49    let mut parts = Vec::new();
50    let mut count = 0;
51    for (chain, numbers) in &by_chain {
52        let numbers = sorted_distinct(numbers);
53        count += numbers.len();
54        let joined = runs(&numbers)
55            .into_iter()
56            .map(|(start, end)| {
57                if start == end {
58                    format!("resid {start}")
59                } else {
60                    format!("resid {start}:{end}")
61                }
62            })
63            .collect::<Vec<_>>()
64            .join(" or ");
65        parts.push(if chain.is_empty() {
66            format!("({joined})")
67        } else {
68            format!("(chain {chain} and ({joined}))")
69        });
70    }
71    (count > 0).then(|| (parts.join(" or "), count))
72}
73
74#[cfg(test)]
75#[path = "residue_ranges_tests.rs"]
76mod tests;
molframe_query/residue_ranges.rs

molframe_query/
residue_ranges.rs
