use crate::error::{Error, Result};
use crate::types::{
AtomFlags, BondDirection, BondFlags, BondOrder, BondStereo, ChiralTag, Hybridization,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AtomData {
pub atomic_num: u8,
pub formal_charge: i8,
pub isotope: u16,
pub num_explicit_hs: u8,
pub num_implicit_hs: u8,
pub num_radical_electrons: u8,
pub atom_map: u16,
pub chiral_tag: ChiralTag,
pub stereo_perm: u8,
pub hybridization: Hybridization,
pub flags: AtomFlags,
}
impl AtomData {
#[must_use]
pub fn new(atomic_num: u8) -> Self {
Self {
atomic_num,
formal_charge: 0,
isotope: 0,
num_explicit_hs: 0,
num_implicit_hs: 0,
num_radical_electrons: 0,
atom_map: 0,
chiral_tag: ChiralTag::Unspecified,
stereo_perm: 0,
hybridization: Hybridization::Unspecified,
flags: AtomFlags::NONE,
}
}
}
impl Default for AtomData {
fn default() -> Self {
Self::new(0)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct BondData {
pub begin: u32,
pub end: u32,
pub order: BondOrder,
pub direction: BondDirection,
pub stereo: BondStereo,
pub stereo_atoms: [u32; 2],
pub flags: BondFlags,
}
impl BondData {
pub const NO_STEREO_ATOM: u32 = u32::MAX;
#[must_use]
pub fn new(begin: u32, end: u32, order: BondOrder) -> Self {
Self {
begin,
end,
order,
direction: BondDirection::None,
stereo: BondStereo::None,
stereo_atoms: [Self::NO_STEREO_ATOM; 2],
flags: BondFlags::NONE,
}
}
#[must_use]
pub fn valence_contribution_to(&self, atom: u32) -> f32 {
if atom != self.begin && atom != self.end {
return 0.0;
}
if self.order == BondOrder::Dative && atom != self.end {
return 0.0; }
self.order.as_double()
}
#[must_use]
pub fn other_end(&self, from: u32) -> Option<u32> {
if from == self.begin {
Some(self.end)
} else if from == self.end {
Some(self.begin)
} else {
None
}
}
}
const NO_HALF: u32 = u32::MAX;
#[derive(Debug, Clone, Default)]
pub struct MolBuilder {
atoms: Vec<AtomData>,
bonds: Vec<BondData>,
name: Option<String>,
first_half: Vec<u32>,
last_half: Vec<u32>,
next_half: Vec<u32>,
degree: Vec<u32>,
}
impl MolBuilder {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_capacity(n_atoms: usize, n_bonds: usize) -> Self {
Self {
atoms: Vec::with_capacity(n_atoms),
bonds: Vec::with_capacity(n_bonds),
name: None,
first_half: Vec::with_capacity(n_atoms),
last_half: Vec::with_capacity(n_atoms),
next_half: Vec::with_capacity(n_bonds * 2),
degree: Vec::with_capacity(n_atoms),
}
}
pub fn add_atom(&mut self, atomic_num: u8) -> u32 {
self.add_atom_data(AtomData::new(atomic_num))
}
pub fn add_atom_data(&mut self, atom: AtomData) -> u32 {
let idx = self.atoms.len() as u32;
self.atoms.push(atom);
self.first_half.push(NO_HALF);
self.last_half.push(NO_HALF);
self.degree.push(0);
idx
}
pub fn add_bond(&mut self, begin: u32, end: u32, order: BondOrder) -> Result<u32> {
self.add_bond_data(BondData::new(begin, end, order))
}
pub fn add_bond_data(&mut self, bond: BondData) -> Result<u32> {
let n = self.atoms.len() as u32;
if bond.begin >= n || bond.end >= n {
return Err(Error::AtomIndexOutOfRange {
index: bond.begin.max(bond.end),
num_atoms: n,
});
}
if bond.begin == bond.end {
return Err(Error::SelfLoop { atom: bond.begin });
}
assert!(
self.bonds.len() < (u32::MAX / 2) as usize,
"键数超出半边编号能表示的范围"
);
let idx = self.bonds.len() as u32;
self.bonds.push(bond);
self.link_half(bond.begin, idx * 2);
self.link_half(bond.end, idx * 2 + 1);
Ok(idx)
}
pub fn swap_bond_ends(&mut self, bond: u32) -> Result<()> {
let num_bonds = self.bonds.len() as u32;
let b = self
.bonds
.get_mut(bond as usize)
.ok_or(Error::BondIndexOutOfRange {
index: bond,
num_bonds,
})?;
std::mem::swap(&mut b.begin, &mut b.end);
self.rebuild_index();
Ok(())
}
fn rebuild_index(&mut self) {
let n = self.atoms.len();
self.first_half.clear();
self.first_half.resize(n, NO_HALF);
self.last_half.clear();
self.last_half.resize(n, NO_HALF);
self.degree.clear();
self.degree.resize(n, 0);
self.next_half.clear();
for i in 0..self.bonds.len() {
let (begin, end) = (self.bonds[i].begin, self.bonds[i].end);
self.link_half(begin, (i * 2) as u32);
self.link_half(end, (i * 2 + 1) as u32);
}
}
fn link_half(&mut self, atom: u32, half: u32) {
debug_assert_eq!(
self.next_half.len() as u32,
half,
"半边必须按编号顺序追加,否则 next_half 的下标语义就断了"
);
self.next_half.push(NO_HALF);
let a = atom as usize;
let tail = self.last_half[a];
if tail == NO_HALF {
self.first_half[a] = half;
} else {
self.next_half[tail as usize] = half;
}
self.last_half[a] = half;
self.degree[a] += 1;
}
#[must_use]
pub fn neighbors(&self, atom: u32) -> Neighbors<'_> {
let head = self
.first_half
.get(atom as usize)
.copied()
.unwrap_or(NO_HALF);
Neighbors {
mol: self,
half: head,
}
}
#[must_use]
pub fn degree(&self, atom: u32) -> usize {
self.degree.get(atom as usize).copied().unwrap_or(0) as usize
}
#[must_use]
pub fn bond_between(&self, a: u32, b: u32) -> Option<u32> {
let from = if self.degree(a) <= self.degree(b) {
a
} else {
b
};
let to = if from == a { b } else { a };
self.neighbors(from)
.find(|&(nbr, _)| nbr == to)
.map(|(_, bi)| bi)
}
#[must_use]
pub fn num_atoms(&self) -> usize {
self.atoms.len()
}
#[must_use]
pub fn num_bonds(&self) -> usize {
self.bonds.len()
}
#[must_use]
pub fn atoms(&self) -> &[AtomData] {
&self.atoms
}
#[must_use]
pub fn bonds(&self) -> &[BondData] {
&self.bonds
}
pub fn atom_mut(&mut self, idx: u32) -> Option<&mut AtomData> {
self.atoms.get_mut(idx as usize)
}
pub fn bond_mut(&mut self, idx: u32) -> Option<BondMut<'_>> {
self.bonds
.get_mut(idx as usize)
.map(|bond| BondMut { bond })
}
#[must_use]
pub fn name(&self) -> Option<&str> {
self.name.as_deref()
}
pub fn set_name(&mut self, name: impl Into<String>) {
self.name = Some(name.into());
}
#[doc(hidden)]
#[must_use]
pub fn adjacency_index_is_consistent(&self) -> bool {
for a in 0..self.atoms.len() as u32 {
let expected: Vec<(u32, u32)> = self
.bonds
.iter()
.enumerate()
.filter_map(|(bi, b)| b.other_end(a).map(|o| (o, bi as u32)))
.collect();
let actual: Vec<(u32, u32)> = self.neighbors(a).collect();
if expected != actual || self.degree(a) != expected.len() {
return false;
}
}
true
}
}
#[derive(Debug, Clone)]
pub struct Neighbors<'a> {
mol: &'a MolBuilder,
half: u32,
}
impl Iterator for Neighbors<'_> {
type Item = (u32, u32);
fn next(&mut self) -> Option<Self::Item> {
if self.half == NO_HALF {
return None;
}
let h = self.half;
self.half = self.mol.next_half[h as usize];
let bi = h >> 1;
let bond = self.mol.bonds[bi as usize];
let nbr = if h & 1 == 0 { bond.end } else { bond.begin };
Some((nbr, bi))
}
}
#[derive(Debug)]
pub struct BondMut<'a> {
bond: &'a mut BondData,
}
impl BondMut<'_> {
#[must_use]
pub fn get(&self) -> BondData {
*self.bond
}
pub fn set_order(&mut self, order: BondOrder) {
self.bond.order = order;
}
pub fn set_direction(&mut self, direction: BondDirection) {
self.bond.direction = direction;
}
pub fn set_stereo(&mut self, stereo: BondStereo) {
self.bond.stereo = stereo;
}
pub fn set_stereo_atoms(&mut self, atoms: [u32; 2]) {
self.bond.stereo_atoms = atoms;
}
pub fn flags_mut(&mut self) -> &mut BondFlags {
&mut self.bond.flags
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn build_ethanol() {
let mut b = MolBuilder::new();
let c0 = b.add_atom(6);
let c1 = b.add_atom(6);
let o = b.add_atom(8);
b.add_bond(c0, c1, BondOrder::Single).unwrap();
b.add_bond(c1, o, BondOrder::Single).unwrap();
assert_eq!(b.num_atoms(), 3);
assert_eq!(b.num_bonds(), 2);
assert_eq!(b.atoms()[2].atomic_num, 8);
}
#[test]
fn rejects_out_of_range_endpoint() {
let mut b = MolBuilder::new();
b.add_atom(6);
let err = b.add_bond(0, 5, BondOrder::Single).unwrap_err();
assert!(matches!(
err,
Error::AtomIndexOutOfRange {
index: 5,
num_atoms: 1
}
));
}
#[test]
fn rejects_self_loop() {
let mut b = MolBuilder::new();
b.add_atom(6);
let err = b.add_bond(0, 0, BondOrder::Single).unwrap_err();
assert!(matches!(err, Error::SelfLoop { atom: 0 }));
}
#[test]
fn bond_other_end() {
let bond = BondData::new(3, 7, BondOrder::Double);
assert_eq!(bond.other_end(3), Some(7));
assert_eq!(bond.other_end(7), Some(3));
assert_eq!(bond.other_end(5), None);
}
}
#[cfg(test)]
mod adjacency_tests {
use super::*;
fn isobutane() -> MolBuilder {
let mut m = MolBuilder::new();
for _ in 0..4 {
m.add_atom(6);
}
m.add_bond(0, 1, BondOrder::Single).unwrap();
m.add_bond(1, 2, BondOrder::Single).unwrap();
m.add_bond(1, 3, BondOrder::Single).unwrap();
m
}
#[test]
fn neighbors_and_degree() {
let m = isobutane();
assert_eq!(
m.neighbors(1).collect::<Vec<_>>(),
vec![(0, 0), (2, 1), (3, 2)]
);
assert_eq!(m.neighbors(0).collect::<Vec<_>>(), vec![(1, 0)]);
assert_eq!(m.degree(1), 3);
assert_eq!(m.degree(0), 1);
}
#[test]
fn isolated_atom_has_no_neighbors() {
let mut m = MolBuilder::new();
m.add_atom(10); assert_eq!(m.neighbors(0).count(), 0);
assert_eq!(m.degree(0), 0);
}
#[test]
fn out_of_range_atom_is_empty_not_panic() {
let m = isobutane();
assert_eq!(m.neighbors(99).count(), 0);
assert_eq!(m.degree(99), 0);
assert_eq!(m.bond_between(99, 0), None);
}
#[test]
fn neighbor_order_is_bond_insertion_order() {
let mut m = MolBuilder::new();
for _ in 0..4 {
m.add_atom(6);
}
m.add_bond(3, 0, BondOrder::Single).unwrap();
m.add_bond(0, 1, BondOrder::Single).unwrap();
m.add_bond(2, 0, BondOrder::Single).unwrap();
assert_eq!(
m.neighbors(0).map(|(a, _)| a).collect::<Vec<_>>(),
vec![3, 1, 2],
"无论中心原子在哪一端,顺序都应是键的插入顺序"
);
assert_eq!(
m.neighbors(0).map(|(_, b)| b).collect::<Vec<_>>(),
vec![0, 1, 2]
);
}
#[test]
fn bond_between_finds_edges_from_either_side() {
let m = isobutane();
assert_eq!(m.bond_between(1, 0), Some(0));
assert_eq!(m.bond_between(0, 1), Some(0));
assert_eq!(m.bond_between(1, 3), Some(2));
assert_eq!(m.bond_between(0, 2), None, "0 与 2 不相邻");
}
#[test]
fn rejected_bond_leaves_index_untouched() {
let mut m = isobutane();
assert!(m.add_bond(1, 1, BondOrder::Single).is_err());
assert!(m.add_bond(0, 99, BondOrder::Single).is_err());
assert_eq!(m.degree(1), 3);
assert_eq!(m.num_bonds(), 3);
assert!(m.adjacency_index_is_consistent());
}
#[test]
fn index_stays_consistent_through_incremental_build() {
let mut m = MolBuilder::new();
assert!(m.adjacency_index_is_consistent());
for i in 0..12u32 {
m.add_atom(6);
assert!(m.adjacency_index_is_consistent(), "加原子 {i} 后失配");
if i > 0 {
m.add_bond(i - 1, i, BondOrder::Single).unwrap();
assert!(m.adjacency_index_is_consistent(), "加键 {i} 后失配");
}
}
m.add_bond(0, 11, BondOrder::Single).unwrap();
m.add_bond(3, 8, BondOrder::Single).unwrap();
assert!(m.adjacency_index_is_consistent());
}
#[test]
fn clone_carries_a_valid_index() {
let m = isobutane().clone();
assert!(m.adjacency_index_is_consistent());
assert_eq!(m.degree(1), 3);
}
#[test]
fn property_edits_do_not_disturb_topology() {
let mut m = isobutane();
let mut b = m.bond_mut(1).unwrap();
b.set_order(BondOrder::Double);
b.flags_mut().insert(BondFlags::AROMATIC);
b.set_direction(BondDirection::UpRight);
assert!(m.adjacency_index_is_consistent());
assert_eq!(m.bonds()[1].order, BondOrder::Double);
assert_eq!(m.bonds()[1].direction, BondDirection::UpRight);
}
}
#[cfg(test)]
mod valence_contrib_tests {
use super::*;
#[test]
fn dative_contribution_is_asymmetric() {
let d = BondData::new(3, 7, BondOrder::Dative);
assert_eq!(d.valence_contribution_to(3), 0.0, "给体不计价");
assert_eq!(d.valence_contribution_to(7), 1.0, "受体计 1");
assert_eq!(d.valence_contribution_to(9), 0.0, "非端点");
}
#[test]
fn normal_bonds_are_symmetric() {
for (order, v) in [
(BondOrder::Single, 1.0),
(BondOrder::Double, 2.0),
(BondOrder::Triple, 3.0),
(BondOrder::Aromatic, 1.5),
] {
let b = BondData::new(1, 2, order);
assert_eq!(b.valence_contribution_to(1), v);
assert_eq!(b.valence_contribution_to(2), v);
}
}
}