xtalforge_core/
selection.rs

1//! Selection model and atom tracking (Task 3.11).
2
3use std::collections::{BTreeMap, BTreeSet};
4
5use serde::{Deserialize, Serialize};
6
7/// A set of atom indices and bond index pairs.
8#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
9pub struct Selection {
10    /// Selected atoms in the current cell.
11    pub atoms: BTreeSet<usize>,
12    /// Selected bonds as ordered index pairs.
13    pub bonds: BTreeSet<(usize, usize)>,
14}
15
16impl Selection {
17    /// Empty selection.
18    #[must_use]
19    pub fn new() -> Self {
20        Self::default()
21    }
22
23    /// Selects one atom.
24    pub fn add_atom(&mut self, index: usize) {
25        self.atoms.insert(index);
26    }
27
28    /// Deselects one atom.
29    pub fn remove_atom(&mut self, index: usize) {
30        self.atoms.remove(&index);
31    }
32
33    /// Toggles one atom.
34    pub fn toggle_atom(&mut self, index: usize) {
35        if !self.atoms.insert(index) {
36            self.atoms.remove(&index);
37        }
38    }
39
40    /// Selects a sorted block of atoms.
41    pub fn add_atoms(&mut self, indices: impl IntoIterator<Item = usize>) {
42        self.atoms.extend(indices);
43    }
44
45    /// Selects one bond; the pair is normalised to `(min, max)`.
46    pub fn add_bond(&mut self, first: usize, second: usize) {
47        self.bonds.insert((first.min(second), first.max(second)));
48    }
49
50    /// Whether an atom is selected.
51    #[must_use]
52    pub fn contains(&self, index: usize) -> bool {
53        self.atoms.contains(&index)
54    }
55
56    /// Number of selected atoms.
57    #[must_use]
58    pub fn len(&self) -> usize {
59        self.atoms.len()
60    }
61
62    /// Whether nothing is selected.
63    #[must_use]
64    pub fn is_empty(&self) -> bool {
65        self.atoms.is_empty() && self.bonds.is_empty()
66    }
67
68    /// Clears atoms and bonds.
69    pub fn clear(&mut self) {
70        self.atoms.clear();
71        self.bonds.clear();
72    }
73
74    /// Union with another selection.
75    pub fn union(&mut self, other: &Self) {
76        self.atoms.extend(other.atoms.iter().copied());
77        self.bonds.extend(other.bonds.iter().copied());
78    }
79
80    /// Intersection with another selection.
81    pub fn intersect(&mut self, other: &Self) {
82        self.atoms.retain(|index| other.atoms.contains(index));
83        self.bonds.retain(|bond| other.bonds.contains(bond));
84    }
85}
86
87/// A stable atom identifier that survives command edits.
88#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
89pub struct AtomId(pub u64);
90
91/// Maps hero atom ids to current site indices.
92#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
93pub struct TrackingTable {
94    /// Current site index by id.
95    pub by_id: BTreeMap<AtomId, usize>,
96    next: u64,
97}
98
99impl TrackingTable {
100    /// Empty table.
101    #[must_use]
102    pub fn new() -> Self {
103        Self::default()
104    }
105
106    /// Assigns a new id to a site and returns it.
107    pub fn assign(&mut self, site: usize) -> AtomId {
108        let id = AtomId(self.next);
109        self.next += 1;
110        self.by_id.insert(id, site);
111        id
112    }
113
114    /// Current site of an id.
115    #[must_use]
116    pub fn site_of(&self, id: AtomId) -> Option<usize> {
117        self.by_id.get(&id).copied()
118    }
119
120    /// Id assigned to a site, if any.
121    #[must_use]
122    pub fn id_of(&self, site: usize) -> Option<AtomId> {
123        self.by_id
124            .iter()
125            .find(|(_, index)| **index == site)
126            .map(|(id, _)| *id)
127    }
128
129    /// Removes an id.
130    pub fn remove(&mut self, id: AtomId) {
131        self.by_id.remove(&id);
132    }
133
134    /// Rewrites every site index through a mapping; missing entries are
135    /// removed.
136    pub fn remap(&mut self, mapping: &BTreeMap<usize, usize>) {
137        self.by_id.retain(|_, site| {
138            if let Some(new) = mapping.get(site) {
139                *site = *new;
140                true
141            } else {
142                false
143            }
144        });
145    }
146}
xtalforge_core/selection.rs

xtalforge_core/
selection.rs
