balsa/tree/
atom.rs

1use crate::feature::{AtomKind, BondKind, Bracket, Selection, Shortcut};
2
3use super::Edge;
4
5#[derive(Debug, PartialEq, Default, Clone)]
6pub struct Atom {
7    pub kind: AtomKind,
8    pub edges: Vec<Edge>,
9}
10
11impl Atom {
12    pub fn new(kind: &AtomKind) -> Self {
13        Self {
14            kind: kind.clone(),
15            edges: Vec::new(),
16        }
17    }
18
19    pub fn star(edges: Vec<Edge>) -> Self {
20        Self {
21            kind: AtomKind::Star,
22            edges,
23        }
24    }
25
26    pub fn shortcut(shortcut: Shortcut, edges: Vec<Edge>) -> Self {
27        Self {
28            kind: AtomKind::Shortcut(shortcut),
29            edges,
30        }
31    }
32
33    pub fn selection(selection: Selection, edges: Vec<Edge>) -> Self {
34        Self {
35            kind: AtomKind::Selection(selection),
36            edges,
37        }
38    }
39
40    pub fn bracket(bracket: Bracket, edges: Vec<Edge>) -> Self {
41        Self {
42            kind: AtomKind::Bracket(bracket),
43            edges,
44        }
45    }
46
47    pub fn valence(&self, input: Option<&BondKind>) -> u8 {
48        let input_bond_order = match input {
49            Some(input) => input.bond_order(),
50            None => 0,
51        };
52        let virtual_hydrogens = self.kind.virtual_hydrogens();
53        let child_bond_order_sum =
54            self.edges.iter().fold(0, |r, e| r + e.bond_order());
55
56        input_bond_order + virtual_hydrogens + child_bond_order_sum
57    }
58
59    pub fn subvalence(&self, input: Option<&BondKind>) -> u8 {
60        self.kind.subvalence(self.valence(input))
61    }
62
63    pub fn implicit_hydrogens(&self, input: Option<&BondKind>) -> u8 {
64        match &self.kind {
65            AtomKind::Star => 0,
66            AtomKind::Shortcut(_) => self.subvalence(input),
67            AtomKind::Selection(_) => match self.subvalence(input) {
68                0 => 0,
69                subvalence => subvalence - 1,
70            },
71            AtomKind::Bracket(_) => 0,
72        }
73    }
74
75    pub fn hydrogens(&self, input: Option<&BondKind>) -> u8 {
76        match &self.kind {
77            AtomKind::Star => 0,
78            AtomKind::Shortcut(_) => self.implicit_hydrogens(input),
79            AtomKind::Selection(_) => self.implicit_hydrogens(input),
80            AtomKind::Bracket(bracket) => bracket.hydrogens(),
81        }
82    }
83}
84
85#[cfg(test)]
86mod valence {
87    use super::*;
88    use crate::feature::VirtualHydrogen;
89    use pretty_assertions::assert_eq;
90
91    #[test]
92    fn zerovalent() {
93        let atom = Atom::star(vec![]);
94        let input = None;
95
96        assert_eq!(atom.valence(input), 0)
97    }
98
99    #[test]
100    fn hydrogens_none_input_double() {
101        let atom = Atom::star(vec![]);
102        let input = Some(BondKind::Double);
103
104        assert_eq!(atom.valence(input.as_ref()), 2)
105    }
106
107    #[test]
108    fn hydrogens_one_input_single() {
109        let atom = Atom::bracket(
110            Bracket {
111                hydrogens: Some(VirtualHydrogen::H),
112                ..Default::default()
113            },
114            vec![],
115        );
116        let input = Some(BondKind::Single);
117
118        assert_eq!(atom.valence(input.as_ref()), 2)
119    }
120
121    #[test]
122    fn hydrogens_one_input_single_children_gap() {
123        let atom = Atom::bracket(
124            Bracket {
125                hydrogens: Some(VirtualHydrogen::H),
126                ..Default::default()
127            },
128            vec![Edge::gap_star(vec![])],
129        );
130        let input = Some(BondKind::Single);
131
132        assert_eq!(atom.valence(input.as_ref()), 2)
133    }
134
135    #[test]
136    fn hydrogens_one_input_signle_children_single() {
137        let atom = Atom::bracket(
138            Bracket {
139                hydrogens: Some(VirtualHydrogen::H),
140                ..Default::default()
141            },
142            vec![Edge::bond_star(BondKind::Double, vec![])],
143        );
144        let input = Some(BondKind::Single);
145
146        assert_eq!(atom.valence(input.as_ref()), 4)
147    }
148}
149
150#[cfg(test)]
151mod subvalence {
152    use crate::feature::{Element, Symbol, VirtualHydrogen};
153
154    use super::*;
155    use pretty_assertions::assert_eq;
156
157    #[test]
158    fn star_zerovalent() {
159        let atom = Atom::star(vec![]);
160        let input = Some(&BondKind::Single);
161
162        assert_eq!(atom.subvalence(input), 0)
163    }
164
165    #[test]
166    fn shortcut_subvalent() {
167        let atom = Atom::shortcut(Shortcut::C, vec![]);
168        let input = Some(&BondKind::Single);
169
170        assert_eq!(atom.subvalence(input), 3)
171    }
172
173    #[test]
174    fn selection_subvalent() {
175        let atom = Atom::selection(Selection::C, vec![]);
176        let input = Some(&BondKind::Single);
177
178        assert_eq!(atom.subvalence(input), 3)
179    }
180
181    #[test]
182    fn bracket_subvalent() {
183        let atom = Atom::bracket(
184            Bracket {
185                symbol: Symbol::Element(Element::C),
186                hydrogens: Some(VirtualHydrogen::H2),
187                ..Default::default()
188            },
189            vec![],
190        );
191        let input = Some(&BondKind::Single);
192
193        assert_eq!(atom.subvalence(input), 1)
194    }
195}
196
197#[cfg(test)]
198mod implicit_hydrogens {
199    use crate::feature;
200    use pretty_assertions::assert_eq;
201
202    use super::*;
203
204    #[test]
205    fn star_zerovalent() {
206        let atom = Atom::star(vec![]);
207        let input = None;
208
209        assert_eq!(atom.implicit_hydrogens(input), 0)
210    }
211
212    #[test]
213    fn shortcut_carbon_input_elided_children_none() {
214        let atom = Atom::shortcut(feature::Shortcut::C, vec![]);
215        let input = Some(&feature::BondKind::Elided);
216
217        assert_eq!(atom.implicit_hydrogens(input), 3)
218    }
219
220    #[test]
221    fn selection_carbon_input_elided_children_none() {
222        let atom = Atom::selection(feature::Selection::C, vec![]);
223        let input = Some(&BondKind::Elided);
224
225        assert_eq!(atom.implicit_hydrogens(input), 2)
226    }
227
228    #[test]
229    fn bracket_carbon_input_elided_children_none() {
230        let atom = Atom::bracket(
231            feature::Bracket {
232                symbol: feature::Symbol::Element(feature::Element::C),
233                ..Default::default()
234            },
235            vec![],
236        );
237        let input = Some(&BondKind::Elided);
238
239        assert_eq!(atom.implicit_hydrogens(input), 0)
240    }
241}
242
243#[cfg(test)]
244mod hydrogens {
245    use super::*;
246    use crate::feature;
247    use pretty_assertions::assert_eq;
248
249    #[test]
250    fn star() {
251        let atom = Atom::star(vec![]);
252        let input = None;
253
254        assert_eq!(atom.hydrogens(input), 0)
255    }
256
257    #[test]
258    fn shortcut() {
259        let atom = Atom::shortcut(feature::Shortcut::C, vec![]);
260        let input = None;
261
262        assert_eq!(atom.hydrogens(input), 4)
263    }
264
265    #[test]
266    fn selection() {
267        let atom = Atom::selection(feature::Selection::C, vec![]);
268        let input = None;
269
270        assert_eq!(atom.hydrogens(input), 3)
271    }
272
273    #[test]
274    fn bracket_hydrogens_some() {
275        let atom = Atom::bracket(
276            feature::Bracket {
277                hydrogens: Some(feature::VirtualHydrogen::H2),
278                ..Default::default()
279            },
280            vec![],
281        );
282        let input = None;
283
284        assert_eq!(atom.hydrogens(input), 2)
285    }
286}