Skip to main content

ak_vis/viewer/session/
bonds.rs

1use std::collections::BTreeSet;
2
3#[derive(Clone, Debug, PartialEq, Eq)]
4pub struct BondList {
5    edges: Vec<(usize, usize)>,
6}
7
8#[derive(Clone, Debug, PartialEq, Eq)]
9pub struct BondFrames {
10    frames: Vec<Option<BondList>>,
11}
12
13impl BondList {
14    pub fn new(edges: impl IntoIterator<Item = (usize, usize)>) -> Self {
15        let mut canonical = BTreeSet::new();
16        for (i, j) in edges {
17            if i == j {
18                continue;
19            }
20            canonical.insert((i.min(j), i.max(j)));
21        }
22        Self {
23            edges: canonical.into_iter().collect(),
24        }
25    }
26
27    pub fn iter(&self) -> impl Iterator<Item = &(usize, usize)> {
28        self.edges.iter()
29    }
30
31    pub fn len(&self) -> usize {
32        self.edges.len()
33    }
34
35    pub fn is_empty(&self) -> bool {
36        self.edges.is_empty()
37    }
38
39    pub fn merged(&self, other: &BondList) -> Self {
40        Self::new(self.iter().copied().chain(other.iter().copied()))
41    }
42
43    pub fn without(&self, other: &BondList) -> Self {
44        let removals: BTreeSet<_> = other.iter().copied().collect();
45        Self::new(self.iter().copied().filter(|edge| !removals.contains(edge)))
46    }
47}
48
49impl BondFrames {
50    pub fn new(frame_count: usize) -> Self {
51        Self {
52            frames: vec![None; frame_count],
53        }
54    }
55
56    pub fn set(&mut self, frame_index: usize, bonds: BondList) {
57        if frame_index < self.frames.len() {
58            self.frames[frame_index] = Some(bonds);
59        }
60    }
61
62    pub fn add(&mut self, frame_index: usize, bonds: BondList) {
63        if frame_index >= self.frames.len() {
64            return;
65        }
66        let merged = match self.frames[frame_index].take() {
67            Some(existing) => existing.merged(&bonds),
68            None => bonds,
69        };
70        self.frames[frame_index] = Some(merged);
71    }
72
73    pub fn remove(&mut self, frame_index: usize, bonds: &BondList) {
74        if frame_index >= self.frames.len() {
75            return;
76        }
77        if let Some(existing) = self.frames[frame_index].take() {
78            let updated = existing.without(bonds);
79            self.frames[frame_index] = (!updated.is_empty()).then_some(updated);
80        }
81    }
82
83    pub fn clear(&mut self, frame_index: usize) {
84        if frame_index < self.frames.len() {
85            self.frames[frame_index] = None;
86        }
87    }
88
89    pub fn get(&self, frame_index: usize) -> Option<&BondList> {
90        self.frames
91            .get(frame_index)
92            .and_then(|bonds| bonds.as_ref())
93    }
94
95    pub fn resize(&mut self, frame_count: usize) {
96        self.frames.resize(frame_count, None);
97    }
98}