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