1use crate::batch::MolBatch;
8use crate::builder::{AtomData, BondData, MolBuilder};
9
10#[derive(Debug, Clone, Copy)]
14pub struct MolView<'a> {
15 batch: &'a MolBatch,
16 idx: u32,
17 atom_base: u32,
18 bond_base: u32,
19 n_atoms: u32,
20 n_bonds: u32,
21}
22
23impl<'a> MolView<'a> {
24 pub(crate) fn new(batch: &'a MolBatch, idx: u32) -> Self {
25 let i = idx as usize;
26 let atom_base = batch.mol_atom_offset[i];
27 let bond_base = batch.mol_bond_offset[i];
28 Self {
29 batch,
30 idx,
31 atom_base,
32 bond_base,
33 n_atoms: batch.mol_atom_offset[i + 1] - atom_base,
34 n_bonds: batch.mol_bond_offset[i + 1] - bond_base,
35 }
36 }
37
38 #[must_use]
40 pub fn index(&self) -> u32 {
41 self.idx
42 }
43
44 #[must_use]
46 pub fn batch(&self) -> &'a MolBatch {
47 self.batch
48 }
49
50 #[must_use]
52 pub fn num_atoms(&self) -> usize {
53 self.n_atoms as usize
54 }
55
56 #[must_use]
58 pub fn num_bonds(&self) -> usize {
59 self.n_bonds as usize
60 }
61
62 #[must_use]
64 pub fn is_empty(&self) -> bool {
65 self.n_atoms == 0
66 }
67
68 #[must_use]
70 pub fn name(&self) -> Option<&'a str> {
71 self.batch.names[self.idx as usize].as_deref()
72 }
73
74 #[must_use]
76 pub fn global_atom(&self, local: u32) -> u32 {
77 self.atom_base + local
78 }
79
80 #[must_use]
82 pub fn local_atom(&self, global: u32) -> Option<u32> {
83 global
84 .checked_sub(self.atom_base)
85 .filter(|&l| l < self.n_atoms)
86 }
87
88 #[must_use]
90 pub fn atom(&self, local: u32) -> Option<AtomData> {
91 if local >= self.n_atoms {
92 return None;
93 }
94 let g = (self.atom_base + local) as usize;
95 let b = self.batch;
96 Some(AtomData {
97 atomic_num: b.atomic_num[g],
98 formal_charge: b.formal_charge[g],
99 isotope: b.isotope[g],
100 num_explicit_hs: b.num_explicit_hs[g],
101 num_implicit_hs: b.num_implicit_hs[g],
102 num_radical_electrons: b.num_radical_electrons[g],
103 atom_map: b.atom_map[g],
104 chiral_tag: b.chiral_tag[g],
105 stereo_perm: b.stereo_perm[g],
106 hybridization: b.hybridization[g],
107 flags: b.atom_flags[g],
108 })
109 }
110
111 #[must_use]
113 pub fn bond(&self, local: u32) -> Option<BondData> {
114 if local >= self.n_bonds {
115 return None;
116 }
117 let g = (self.bond_base + local) as usize;
118 let b = self.batch;
119 Some(BondData {
120 begin: b.bond_begin[g] - self.atom_base,
121 end: b.bond_end[g] - self.atom_base,
122 order: b.bond_order[g],
123 direction: b.bond_direction[g],
124 stereo: b.bond_stereo[g],
125 stereo_atoms: b.bond_stereo_atoms[g].map(|a| {
126 if a == BondData::NO_STEREO_ATOM {
127 a
128 } else {
129 a - self.atom_base
130 }
131 }),
132 flags: b.bond_flags[g],
133 })
134 }
135
136 #[must_use]
138 pub fn degree(&self, local: u32) -> usize {
139 if local >= self.n_atoms {
140 return 0;
141 }
142 let g = (self.atom_base + local) as usize;
143 (self.batch.nbr_offset[g + 1] - self.batch.nbr_offset[g]) as usize
144 }
145
146 pub fn neighbors(&self, local: u32) -> impl Iterator<Item = (u32, u32)> + 'a {
151 let (start, end) = if local < self.n_atoms {
152 let g = (self.atom_base + local) as usize;
153 (self.batch.nbr_offset[g], self.batch.nbr_offset[g + 1])
154 } else {
155 (0, 0)
156 };
157 let b = self.batch;
158 let (abase, bbase) = (self.atom_base, self.bond_base);
159 (start..end).map(move |k| {
160 let k = k as usize;
161 (b.nbr_atom[k] - abase, b.nbr_bond[k] - bbase)
162 })
163 }
164
165 pub fn atoms(&self) -> impl Iterator<Item = (u32, AtomData)> + '_ {
167 (0..self.n_atoms).map(move |i| (i, self.atom(i).expect("下标由 n_atoms 生成")))
168 }
169
170 pub fn bonds(&self) -> impl Iterator<Item = (u32, BondData)> + '_ {
172 (0..self.n_bonds).map(move |i| (i, self.bond(i).expect("下标由 n_bonds 生成")))
173 }
174
175 #[must_use]
179 pub fn to_builder(&self) -> MolBuilder {
180 let mut m = MolBuilder::with_capacity(self.num_atoms(), self.num_bonds());
181 for (_, a) in self.atoms() {
182 m.add_atom_data(a);
183 }
184 for (_, bd) in self.bonds() {
185 m.add_bond_data(bd).expect("视图中的键端点必然合法");
186 }
187 if let Some(n) = self.name() {
188 m.set_name(n);
189 }
190 m
191 }
192}
193
194#[cfg(test)]
195mod tests {
196 use super::*;
197 use crate::batch::MolBatchBuilder;
198 use crate::types::{
199 AtomFlags, BondDirection, BondFlags, BondOrder, BondStereo, ChiralTag, Hybridization,
200 };
201
202 fn ethanol() -> MolBuilder {
203 let mut m = MolBuilder::new();
204 let c0 = m.add_atom(6);
205 let c1 = m.add_atom(6);
206 let o = m.add_atom(8);
207 m.add_bond(c0, c1, BondOrder::Single).unwrap();
208 m.add_bond(c1, o, BondOrder::Single).unwrap();
209 m.set_name("ethanol");
210 m
211 }
212
213 fn acetic_acid() -> MolBuilder {
214 let mut m = MolBuilder::new();
216 let c0 = m.add_atom(6);
217 let c1 = m.add_atom(6);
218 let o1 = m.add_atom(8);
219 let o2 = m.add_atom(8);
220 m.add_bond(c0, c1, BondOrder::Single).unwrap();
221 m.add_bond(c1, o1, BondOrder::Double).unwrap();
222 m.add_bond(c1, o2, BondOrder::Single).unwrap();
223 m
224 }
225
226 fn batch_of(mols: &[MolBuilder]) -> MolBatch {
227 let mut bb = MolBatchBuilder::new();
228 for m in mols {
229 bb.push(m).unwrap();
230 }
231 bb.finish()
232 }
233
234 #[test]
235 fn local_indices_are_zero_based_per_molecule() {
236 let b = batch_of(&[ethanol(), acetic_acid()]);
237 let m1 = b.mol(1).unwrap();
238
239 assert_eq!(m1.num_atoms(), 4);
240 assert_eq!(m1.num_bonds(), 3);
241 assert_eq!(m1.global_atom(0), 3);
243 assert_eq!(m1.local_atom(3), Some(0));
244 assert_eq!(m1.local_atom(2), None, "全局原子 2 属于第一个分子");
245 assert_eq!(m1.local_atom(7), None, "全局原子 7 已越过本分子");
246
247 let bond = m1.bond(1).unwrap();
248 assert_eq!((bond.begin, bond.end), (1, 2), "键端点应为局部下标");
249 assert_eq!(bond.order, BondOrder::Double);
250 }
251
252 #[test]
253 fn neighbors_are_local() {
254 let b = batch_of(&[ethanol(), acetic_acid()]);
255 let m1 = b.mol(1).unwrap();
256
257 let mut nbrs: Vec<u32> = m1.neighbors(1).map(|(a, _)| a).collect();
259 nbrs.sort_unstable();
260 assert_eq!(nbrs, vec![0, 2, 3]);
261 assert_eq!(m1.degree(1), 3);
262 assert_eq!(m1.degree(0), 1);
263 }
264
265 #[test]
266 fn neighbor_bond_indices_are_local() {
267 let b = batch_of(&[ethanol(), acetic_acid()]);
268 let m1 = b.mol(1).unwrap();
269 for (_, bond_local) in m1.neighbors(1) {
270 assert!(
271 (bond_local as usize) < m1.num_bonds(),
272 "键下标 {bond_local} 未换算为局部"
273 );
274 }
275 }
276
277 #[test]
278 fn out_of_range_access_is_none() {
279 let b = batch_of(&[ethanol()]);
280 let m = b.mol(0).unwrap();
281 assert!(m.atom(3).is_none());
282 assert!(m.bond(2).is_none());
283 assert_eq!(m.degree(99), 0);
284 assert_eq!(m.neighbors(99).count(), 0);
285 }
286
287 fn every_field_set() -> MolBuilder {
297 let mut m = MolBuilder::new();
298 m.add_atom_data(AtomData {
299 atomic_num: 7,
300 formal_charge: -1,
301 isotope: 15,
302 num_explicit_hs: 2,
303 num_implicit_hs: 3,
304 num_radical_electrons: 1,
305 atom_map: 7,
306 chiral_tag: ChiralTag::Cw,
307 stereo_perm: 0,
308 hybridization: Hybridization::Sp3d2,
309 flags: AtomFlags::AROMATIC | AtomFlags::NO_IMPLICIT | AtomFlags::IN_RING,
310 });
311 m.add_atom_data(AtomData {
312 atomic_num: 16,
313 formal_charge: 2,
314 isotope: 34,
315 num_explicit_hs: 1,
316 num_implicit_hs: 0,
317 num_radical_electrons: 2,
318 atom_map: 3,
319 chiral_tag: ChiralTag::Octahedral,
320 stereo_perm: 25,
321 hybridization: Hybridization::Sp2d,
322 flags: AtomFlags::CONJUGATED,
323 });
324 m.add_bond_data(BondData {
325 begin: 0,
326 end: 1,
327 order: BondOrder::Dative,
328 direction: BondDirection::DownRight,
329 stereo: BondStereo::Trans,
330 stereo_atoms: [1, 0],
333 flags: BondFlags::AROMATIC | BondFlags::IN_RING | BondFlags::CONJUGATED,
334 })
335 .expect("端点合法");
336 m.set_name("每字段非默认");
337 m
338 }
339
340 #[test]
343 fn roundtrip_through_batch_is_identity() {
344 let m = every_field_set();
345
346 let b = batch_of(&[ethanol(), m.clone(), ethanol()]);
348 let back = b.mol(1).unwrap().to_builder();
349
350 assert_eq!(back.atoms(), m.atoms(), "原子列往返不一致");
351 assert_eq!(back.bonds(), m.bonds(), "键列往返不一致");
352 assert_eq!(back.name(), m.name());
353 }
354
355 #[test]
358 fn every_field_is_actually_non_default() {
359 let m = every_field_set();
360 let default_atom = AtomData::default();
361 for (i, a) in m.atoms().iter().enumerate() {
362 assert_ne!(a.atomic_num, default_atom.atomic_num, "原子{i}.元素");
363 assert_ne!(a.formal_charge, default_atom.formal_charge, "原子{i}.电荷");
364 assert_ne!(a.isotope, default_atom.isotope, "原子{i}.同位素");
365 assert_ne!(a.atom_map, default_atom.atom_map, "原子{i}.映射号");
366 assert_ne!(a.chiral_tag, default_atom.chiral_tag, "原子{i}.手性");
367 assert_ne!(a.hybridization, default_atom.hybridization, "原子{i}.杂化");
368 assert_ne!(a.flags, AtomFlags::NONE, "原子{i}.标志");
369 }
370 assert!(m.atoms().iter().any(|a| a.num_explicit_hs != 0));
372 assert!(m.atoms().iter().any(|a| a.num_implicit_hs != 0));
373 assert!(m.atoms().iter().any(|a| a.num_radical_electrons != 0));
374 assert!(m.atoms().iter().any(|a| a.stereo_perm != 0));
377 assert!(m.atoms().iter().any(|a| a.chiral_tag.is_tetrahedral()));
378 assert!(m.atoms().iter().any(|a| !a.chiral_tag.is_tetrahedral()));
379
380 for (i, b) in m.bonds().iter().enumerate() {
381 assert_ne!(b.order, BondOrder::Unspecified, "键{i}.键级");
382 assert_ne!(b.direction, BondDirection::None, "键{i}.方向");
383 assert_ne!(b.stereo, BondStereo::None, "键{i}.立体");
384 assert_eq!(b.stereo_atoms, [1, 0], "键{i}.立体参照原子");
387 assert_ne!(b.flags, BondFlags::NONE, "键{i}.标志");
388 }
389 }
390
391 #[test]
392 fn iter_visits_every_molecule() {
393 let b = batch_of(&[ethanol(), acetic_acid(), ethanol()]);
394 let sizes: Vec<usize> = b.iter().map(|m| m.num_atoms()).collect();
395 assert_eq!(sizes, vec![3, 4, 3]);
396 assert_eq!(b.iter().count(), 3);
397 }
398
399 #[test]
400 fn try_mol_reports_out_of_range() {
401 let b = batch_of(&[ethanol()]);
402 assert!(b.try_mol(0).is_ok());
403 let e = b.try_mol(5).unwrap_err();
404 assert!(matches!(
405 e,
406 crate::error::Error::MolIndexOutOfRange {
407 index: 5,
408 num_mols: 1
409 }
410 ));
411 }
412}