1use crate::error::{Error, Result};
46use crate::types::{
47 AtomFlags, BondDirection, BondFlags, BondOrder, BondStereo, ChiralTag, Hybridization,
48};
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub struct AtomData {
55 pub atomic_num: u8,
57 pub formal_charge: i8,
59 pub isotope: u16,
61 pub num_explicit_hs: u8,
63 pub num_implicit_hs: u8,
65 pub num_radical_electrons: u8,
70 pub atom_map: u16,
72 pub chiral_tag: ChiralTag,
74 pub stereo_perm: u8,
109 pub hybridization: Hybridization,
111 pub flags: AtomFlags,
113}
114
115impl AtomData {
116 #[must_use]
118 pub fn new(atomic_num: u8) -> Self {
119 Self {
120 atomic_num,
121 formal_charge: 0,
122 isotope: 0,
123 num_explicit_hs: 0,
124 num_implicit_hs: 0,
125 num_radical_electrons: 0,
126 atom_map: 0,
127 chiral_tag: ChiralTag::Unspecified,
128 stereo_perm: 0,
129 hybridization: Hybridization::Unspecified,
130 flags: AtomFlags::NONE,
131 }
132 }
133}
134
135impl Default for AtomData {
136 fn default() -> Self {
137 Self::new(0)
138 }
139}
140
141#[derive(Debug, Clone, Copy, PartialEq, Eq)]
143pub struct BondData {
144 pub begin: u32,
146 pub end: u32,
148 pub order: BondOrder,
150 pub direction: BondDirection,
152 pub stereo: BondStereo,
155 pub stereo_atoms: [u32; 2],
169 pub flags: BondFlags,
171}
172
173impl BondData {
174 pub const NO_STEREO_ATOM: u32 = u32::MAX;
176
177 #[must_use]
179 pub fn new(begin: u32, end: u32, order: BondOrder) -> Self {
180 Self {
181 begin,
182 end,
183 order,
184 direction: BondDirection::None,
185 stereo: BondStereo::None,
186 stereo_atoms: [Self::NO_STEREO_ATOM; 2],
187 flags: BondFlags::NONE,
188 }
189 }
190
191 #[must_use]
199 pub fn valence_contribution_to(&self, atom: u32) -> f32 {
200 if atom != self.begin && atom != self.end {
201 return 0.0;
202 }
203 if self.order == BondOrder::Dative && atom != self.end {
204 return 0.0; }
206 self.order.as_double()
207 }
208
209 #[must_use]
211 pub fn other_end(&self, from: u32) -> Option<u32> {
212 if from == self.begin {
213 Some(self.end)
214 } else if from == self.end {
215 Some(self.begin)
216 } else {
217 None
218 }
219 }
220}
221
222const NO_HALF: u32 = u32::MAX;
224
225#[derive(Debug, Clone, Default)]
229pub struct MolBuilder {
230 atoms: Vec<AtomData>,
231 bonds: Vec<BondData>,
232 name: Option<String>,
233
234 first_half: Vec<u32>,
241 last_half: Vec<u32>,
243 next_half: Vec<u32>,
245 degree: Vec<u32>,
247}
248
249impl MolBuilder {
250 #[must_use]
252 pub fn new() -> Self {
253 Self::default()
254 }
255
256 #[must_use]
258 pub fn with_capacity(n_atoms: usize, n_bonds: usize) -> Self {
259 Self {
260 atoms: Vec::with_capacity(n_atoms),
261 bonds: Vec::with_capacity(n_bonds),
262 name: None,
263 first_half: Vec::with_capacity(n_atoms),
264 last_half: Vec::with_capacity(n_atoms),
265 next_half: Vec::with_capacity(n_bonds * 2),
266 degree: Vec::with_capacity(n_atoms),
267 }
268 }
269
270 pub fn add_atom(&mut self, atomic_num: u8) -> u32 {
272 self.add_atom_data(AtomData::new(atomic_num))
273 }
274
275 pub fn add_atom_data(&mut self, atom: AtomData) -> u32 {
277 let idx = self.atoms.len() as u32;
278 self.atoms.push(atom);
279 self.first_half.push(NO_HALF);
280 self.last_half.push(NO_HALF);
281 self.degree.push(0);
282 idx
283 }
284
285 pub fn add_bond(&mut self, begin: u32, end: u32, order: BondOrder) -> Result<u32> {
291 self.add_bond_data(BondData::new(begin, end, order))
292 }
293
294 pub fn add_bond_data(&mut self, bond: BondData) -> Result<u32> {
305 let n = self.atoms.len() as u32;
306 if bond.begin >= n || bond.end >= n {
307 return Err(Error::AtomIndexOutOfRange {
308 index: bond.begin.max(bond.end),
309 num_atoms: n,
310 });
311 }
312 if bond.begin == bond.end {
313 return Err(Error::SelfLoop { atom: bond.begin });
314 }
315 assert!(
316 self.bonds.len() < (u32::MAX / 2) as usize,
317 "键数超出半边编号能表示的范围"
318 );
319 let idx = self.bonds.len() as u32;
320 self.bonds.push(bond);
321 self.link_half(bond.begin, idx * 2);
322 self.link_half(bond.end, idx * 2 + 1);
323 Ok(idx)
324 }
325
326 pub fn swap_bond_ends(&mut self, bond: u32) -> Result<()> {
344 let num_bonds = self.bonds.len() as u32;
345 let b = self
346 .bonds
347 .get_mut(bond as usize)
348 .ok_or(Error::BondIndexOutOfRange {
349 index: bond,
350 num_bonds,
351 })?;
352 std::mem::swap(&mut b.begin, &mut b.end);
353 self.rebuild_index();
354 Ok(())
355 }
356
357 fn rebuild_index(&mut self) {
359 let n = self.atoms.len();
360 self.first_half.clear();
361 self.first_half.resize(n, NO_HALF);
362 self.last_half.clear();
363 self.last_half.resize(n, NO_HALF);
364 self.degree.clear();
365 self.degree.resize(n, 0);
366 self.next_half.clear();
367
368 for i in 0..self.bonds.len() {
369 let (begin, end) = (self.bonds[i].begin, self.bonds[i].end);
370 self.link_half(begin, (i * 2) as u32);
371 self.link_half(end, (i * 2 + 1) as u32);
372 }
373 }
374
375 fn link_half(&mut self, atom: u32, half: u32) {
379 debug_assert_eq!(
380 self.next_half.len() as u32,
381 half,
382 "半边必须按编号顺序追加,否则 next_half 的下标语义就断了"
383 );
384 self.next_half.push(NO_HALF);
385
386 let a = atom as usize;
387 let tail = self.last_half[a];
388 if tail == NO_HALF {
389 self.first_half[a] = half;
390 } else {
391 self.next_half[tail as usize] = half;
392 }
393 self.last_half[a] = half;
394 self.degree[a] += 1;
395 }
396
397 #[must_use]
404 pub fn neighbors(&self, atom: u32) -> Neighbors<'_> {
405 let head = self
406 .first_half
407 .get(atom as usize)
408 .copied()
409 .unwrap_or(NO_HALF);
410 Neighbors {
411 mol: self,
412 half: head,
413 }
414 }
415
416 #[must_use]
418 pub fn degree(&self, atom: u32) -> usize {
419 self.degree.get(atom as usize).copied().unwrap_or(0) as usize
420 }
421
422 #[must_use]
424 pub fn bond_between(&self, a: u32, b: u32) -> Option<u32> {
425 let from = if self.degree(a) <= self.degree(b) {
427 a
428 } else {
429 b
430 };
431 let to = if from == a { b } else { a };
432 self.neighbors(from)
433 .find(|&(nbr, _)| nbr == to)
434 .map(|(_, bi)| bi)
435 }
436
437 #[must_use]
439 pub fn num_atoms(&self) -> usize {
440 self.atoms.len()
441 }
442
443 #[must_use]
445 pub fn num_bonds(&self) -> usize {
446 self.bonds.len()
447 }
448
449 #[must_use]
451 pub fn atoms(&self) -> &[AtomData] {
452 &self.atoms
453 }
454
455 #[must_use]
457 pub fn bonds(&self) -> &[BondData] {
458 &self.bonds
459 }
460
461 pub fn atom_mut(&mut self, idx: u32) -> Option<&mut AtomData> {
463 self.atoms.get_mut(idx as usize)
464 }
465
466 pub fn bond_mut(&mut self, idx: u32) -> Option<BondMut<'_>> {
472 self.bonds
473 .get_mut(idx as usize)
474 .map(|bond| BondMut { bond })
475 }
476
477 #[must_use]
479 pub fn name(&self) -> Option<&str> {
480 self.name.as_deref()
481 }
482
483 pub fn set_name(&mut self, name: impl Into<String>) {
485 self.name = Some(name.into());
486 }
487
488 #[doc(hidden)]
493 #[must_use]
494 pub fn adjacency_index_is_consistent(&self) -> bool {
495 for a in 0..self.atoms.len() as u32 {
496 let expected: Vec<(u32, u32)> = self
497 .bonds
498 .iter()
499 .enumerate()
500 .filter_map(|(bi, b)| b.other_end(a).map(|o| (o, bi as u32)))
501 .collect();
502 let actual: Vec<(u32, u32)> = self.neighbors(a).collect();
503 if expected != actual || self.degree(a) != expected.len() {
504 return false;
505 }
506 }
507 true
508 }
509}
510
511#[derive(Debug, Clone)]
513pub struct Neighbors<'a> {
514 mol: &'a MolBuilder,
515 half: u32,
516}
517
518impl Iterator for Neighbors<'_> {
519 type Item = (u32, u32);
521
522 fn next(&mut self) -> Option<Self::Item> {
523 if self.half == NO_HALF {
524 return None;
525 }
526 let h = self.half;
527 self.half = self.mol.next_half[h as usize];
528
529 let bi = h >> 1;
530 let bond = self.mol.bonds[bi as usize];
531 let nbr = if h & 1 == 0 { bond.end } else { bond.begin };
533 Some((nbr, bi))
534 }
535}
536
537#[derive(Debug)]
541pub struct BondMut<'a> {
542 bond: &'a mut BondData,
543}
544
545impl BondMut<'_> {
546 #[must_use]
548 pub fn get(&self) -> BondData {
549 *self.bond
550 }
551
552 pub fn set_order(&mut self, order: BondOrder) {
554 self.bond.order = order;
555 }
556
557 pub fn set_direction(&mut self, direction: BondDirection) {
559 self.bond.direction = direction;
560 }
561
562 pub fn set_stereo(&mut self, stereo: BondStereo) {
564 self.bond.stereo = stereo;
565 }
566
567 pub fn set_stereo_atoms(&mut self, atoms: [u32; 2]) {
572 self.bond.stereo_atoms = atoms;
573 }
574
575 pub fn flags_mut(&mut self) -> &mut BondFlags {
577 &mut self.bond.flags
578 }
579}
580
581#[cfg(test)]
582mod tests {
583 use super::*;
584
585 #[test]
586 fn build_ethanol() {
587 let mut b = MolBuilder::new();
588 let c0 = b.add_atom(6);
589 let c1 = b.add_atom(6);
590 let o = b.add_atom(8);
591 b.add_bond(c0, c1, BondOrder::Single).unwrap();
592 b.add_bond(c1, o, BondOrder::Single).unwrap();
593
594 assert_eq!(b.num_atoms(), 3);
595 assert_eq!(b.num_bonds(), 2);
596 assert_eq!(b.atoms()[2].atomic_num, 8);
597 }
598
599 #[test]
600 fn rejects_out_of_range_endpoint() {
601 let mut b = MolBuilder::new();
602 b.add_atom(6);
603 let err = b.add_bond(0, 5, BondOrder::Single).unwrap_err();
604 assert!(matches!(
605 err,
606 Error::AtomIndexOutOfRange {
607 index: 5,
608 num_atoms: 1
609 }
610 ));
611 }
612
613 #[test]
614 fn rejects_self_loop() {
615 let mut b = MolBuilder::new();
616 b.add_atom(6);
617 let err = b.add_bond(0, 0, BondOrder::Single).unwrap_err();
618 assert!(matches!(err, Error::SelfLoop { atom: 0 }));
619 }
620
621 #[test]
622 fn bond_other_end() {
623 let bond = BondData::new(3, 7, BondOrder::Double);
624 assert_eq!(bond.other_end(3), Some(7));
625 assert_eq!(bond.other_end(7), Some(3));
626 assert_eq!(bond.other_end(5), None);
627 }
628}
629
630#[cfg(test)]
631mod adjacency_tests {
632 use super::*;
633
634 fn isobutane() -> MolBuilder {
636 let mut m = MolBuilder::new();
637 for _ in 0..4 {
638 m.add_atom(6);
639 }
640 m.add_bond(0, 1, BondOrder::Single).unwrap();
641 m.add_bond(1, 2, BondOrder::Single).unwrap();
642 m.add_bond(1, 3, BondOrder::Single).unwrap();
643 m
644 }
645
646 #[test]
647 fn neighbors_and_degree() {
648 let m = isobutane();
649 assert_eq!(
650 m.neighbors(1).collect::<Vec<_>>(),
651 vec![(0, 0), (2, 1), (3, 2)]
652 );
653 assert_eq!(m.neighbors(0).collect::<Vec<_>>(), vec![(1, 0)]);
654 assert_eq!(m.degree(1), 3);
655 assert_eq!(m.degree(0), 1);
656 }
657
658 #[test]
659 fn isolated_atom_has_no_neighbors() {
660 let mut m = MolBuilder::new();
661 m.add_atom(10); assert_eq!(m.neighbors(0).count(), 0);
663 assert_eq!(m.degree(0), 0);
664 }
665
666 #[test]
667 fn out_of_range_atom_is_empty_not_panic() {
668 let m = isobutane();
669 assert_eq!(m.neighbors(99).count(), 0);
670 assert_eq!(m.degree(99), 0);
671 assert_eq!(m.bond_between(99, 0), None);
672 }
673
674 #[test]
677 fn neighbor_order_is_bond_insertion_order() {
678 let mut m = MolBuilder::new();
679 for _ in 0..4 {
680 m.add_atom(6);
681 }
682 m.add_bond(3, 0, BondOrder::Single).unwrap();
684 m.add_bond(0, 1, BondOrder::Single).unwrap();
685 m.add_bond(2, 0, BondOrder::Single).unwrap();
686
687 assert_eq!(
688 m.neighbors(0).map(|(a, _)| a).collect::<Vec<_>>(),
689 vec![3, 1, 2],
690 "无论中心原子在哪一端,顺序都应是键的插入顺序"
691 );
692 assert_eq!(
693 m.neighbors(0).map(|(_, b)| b).collect::<Vec<_>>(),
694 vec![0, 1, 2]
695 );
696 }
697
698 #[test]
699 fn bond_between_finds_edges_from_either_side() {
700 let m = isobutane();
701 assert_eq!(m.bond_between(1, 0), Some(0));
702 assert_eq!(m.bond_between(0, 1), Some(0));
703 assert_eq!(m.bond_between(1, 3), Some(2));
704 assert_eq!(m.bond_between(0, 2), None, "0 与 2 不相邻");
705 }
706
707 #[test]
710 fn rejected_bond_leaves_index_untouched() {
711 let mut m = isobutane();
712 assert!(m.add_bond(1, 1, BondOrder::Single).is_err());
713 assert!(m.add_bond(0, 99, BondOrder::Single).is_err());
714 assert_eq!(m.degree(1), 3);
715 assert_eq!(m.num_bonds(), 3);
716 assert!(m.adjacency_index_is_consistent());
717 }
718
719 #[test]
721 fn index_stays_consistent_through_incremental_build() {
722 let mut m = MolBuilder::new();
723 assert!(m.adjacency_index_is_consistent());
724 for i in 0..12u32 {
725 m.add_atom(6);
726 assert!(m.adjacency_index_is_consistent(), "加原子 {i} 后失配");
727 if i > 0 {
728 m.add_bond(i - 1, i, BondOrder::Single).unwrap();
729 assert!(m.adjacency_index_is_consistent(), "加键 {i} 后失配");
730 }
731 }
732 m.add_bond(0, 11, BondOrder::Single).unwrap();
734 m.add_bond(3, 8, BondOrder::Single).unwrap();
735 assert!(m.adjacency_index_is_consistent());
736 }
737
738 #[test]
740 fn clone_carries_a_valid_index() {
741 let m = isobutane().clone();
742 assert!(m.adjacency_index_is_consistent());
743 assert_eq!(m.degree(1), 3);
744 }
745
746 #[test]
748 fn property_edits_do_not_disturb_topology() {
749 let mut m = isobutane();
750 let mut b = m.bond_mut(1).unwrap();
751 b.set_order(BondOrder::Double);
752 b.flags_mut().insert(BondFlags::AROMATIC);
753 b.set_direction(BondDirection::UpRight);
754 assert!(m.adjacency_index_is_consistent());
755 assert_eq!(m.bonds()[1].order, BondOrder::Double);
756 assert_eq!(m.bonds()[1].direction, BondDirection::UpRight);
757 }
758}
759
760#[cfg(test)]
761mod valence_contrib_tests {
762 use super::*;
763
764 #[test]
766 fn dative_contribution_is_asymmetric() {
767 let d = BondData::new(3, 7, BondOrder::Dative);
768 assert_eq!(d.valence_contribution_to(3), 0.0, "给体不计价");
769 assert_eq!(d.valence_contribution_to(7), 1.0, "受体计 1");
770 assert_eq!(d.valence_contribution_to(9), 0.0, "非端点");
771 }
772
773 #[test]
774 fn normal_bonds_are_symmetric() {
775 for (order, v) in [
776 (BondOrder::Single, 1.0),
777 (BondOrder::Double, 2.0),
778 (BondOrder::Triple, 3.0),
779 (BondOrder::Aromatic, 1.5),
780 ] {
781 let b = BondData::new(1, 2, order);
782 assert_eq!(b.valence_contribution_to(1), v);
783 assert_eq!(b.valence_contribution_to(2), v);
784 }
785 }
786}