chematic_core/

molecule.rs

//! Molecule graph: atoms, bonds, and adjacency list.

use crate::atom::Atom;
use crate::bond::{BondEntry, BondOrder};
use crate::element::Element;

/// Newtype index for an atom in a Molecule.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct AtomIdx(pub u32);

/// Newtype index for a bond in a Molecule.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct BondIdx(pub u32);

/// Error types for molecule construction.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MolError {
    /// Atom index out of range.
    InvalidAtomIdx(AtomIdx),
    /// Duplicate bond between the same pair of atoms.
    DuplicateBond(AtomIdx, AtomIdx),
}

impl core::fmt::Display for MolError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::InvalidAtomIdx(idx) => write!(f, "invalid atom index: {}", idx.0),
            Self::DuplicateBond(a, b) => write!(f, "duplicate bond between atoms {} and {}", a.0, b.0),
        }
    }
}

/// An immutable molecular graph built via [`MoleculeBuilder`].
///
/// Representation: atom list + bond list + per-atom adjacency list.
/// No external graph library is used; all graph traversal is domain-aware.
pub struct Molecule {
    atoms: Vec<Atom>,
    bonds: Vec<BondEntry>,
    /// adjacency[atom_idx] = list of (neighbor_atom_idx, bond_idx)
    adjacency: Vec<Vec<(AtomIdx, BondIdx)>>,
}

impl Molecule {
    /// Number of heavy atoms (does not count implicit H).
    pub fn atom_count(&self) -> usize {
        self.atoms.len()
    }

    /// Number of bonds (edges).
    pub fn bond_count(&self) -> usize {
        self.bonds.len()
    }

    /// Borrow atom by index.
    ///
    /// # Panics
    /// Panics if `idx` is out of range (should not happen with indices from this molecule).
    pub fn atom(&self, idx: AtomIdx) -> &Atom {
        &self.atoms[idx.0 as usize]
    }

    /// Borrow bond by index.
    pub fn bond(&self, idx: BondIdx) -> &BondEntry {
        &self.bonds[idx.0 as usize]
    }

    /// Iterate over all atoms as `(AtomIdx, &Atom)`.
    pub fn atoms(&self) -> impl Iterator<Item = (AtomIdx, &Atom)> {
        self.atoms
            .iter()
            .enumerate()
            .map(|(i, a)| (AtomIdx(i as u32), a))
    }

    /// Iterate over all bonds as `(BondIdx, &BondEntry)`.
    pub fn bonds(&self) -> impl Iterator<Item = (BondIdx, &BondEntry)> {
        self.bonds
            .iter()
            .enumerate()
            .map(|(i, b)| (BondIdx(i as u32), b))
    }

    /// Iterate over neighbors of `idx` as `(neighbor_atom_idx, bond_idx)`.
    pub fn neighbors(&self, idx: AtomIdx) -> impl Iterator<Item = (AtomIdx, BondIdx)> + '_ {
        self.adjacency[idx.0 as usize].iter().copied()
    }

    /// Degree (number of connected bonds) of atom `idx`.
    pub fn degree(&self, idx: AtomIdx) -> usize {
        self.adjacency[idx.0 as usize].len()
    }

    /// Return the bond between `a` and `b`, or `None` if not connected.
    pub fn bond_between(&self, a: AtomIdx, b: AtomIdx) -> Option<(BondIdx, &BondEntry)> {
        self.adjacency[a.0 as usize]
            .iter()
            .find(|&&(nb, _)| nb == b)
            .map(|&(_, bidx)| (bidx, &self.bonds[bidx.0 as usize]))
    }

    /// Molecular formula as a Hill-order string (C first, H second, then alphabetical).
    pub fn formula(&self) -> String {
        use std::collections::BTreeMap;
        let mut counts: BTreeMap<&str, u32> = BTreeMap::new();
        for (_, atom) in self.atoms() {
            *counts.entry(atom.element.symbol()).or_insert(0) += 1;
        }

        let mut result = String::new();
        let push_count = |sym: &str, n: u32, out: &mut String| {
            out.push_str(sym);
            if n > 1 {
                out.push_str(&n.to_string());
            }
        };

        // Hill order: C, then H, then the remaining symbols alphabetically.
        if let Some(c) = counts.remove("C") {
            push_count("C", c, &mut result);
        }
        if let Some(h) = counts.remove("H") {
            push_count("H", h, &mut result);
        }
        for (sym, count) in &counts {
            push_count(sym, *count, &mut result);
        }
        result
    }
}

// ---------------------------------------------------------------------------
// Immutable update methods (functional-style editing)
// ---------------------------------------------------------------------------

impl Molecule {
    /// Return a new `Molecule` with one extra atom appended, along with the
    /// index that the new atom will have in the returned molecule.
    pub fn with_atom_added(&self, atom: Atom) -> (Molecule, AtomIdx) {
        let mut builder = MoleculeBuilder::new();
        for (_, a) in self.atoms() {
            builder.add_atom(a.clone());
        }
        for (_, b) in self.bonds() {
            let _ = builder.add_bond(b.atom1, b.atom2, b.order);
        }
        let new_idx = builder.add_atom(atom);
        (builder.build(), new_idx)
    }

    /// Return a new `Molecule` with one extra bond added, along with the index
    /// of the newly added bond in the returned molecule.
    ///
    /// Returns `Err` if `a == b` or the bond already exists (same semantics as
    /// [`MoleculeBuilder::add_bond`]).
    pub fn with_bond_added(
        &self,
        a: AtomIdx,
        b: AtomIdx,
        order: BondOrder,
    ) -> Result<(Molecule, BondIdx), MolError> {
        let mut builder = MoleculeBuilder::new();
        for (_, atom) in self.atoms() {
            builder.add_atom(atom.clone());
        }
        for (_, bond) in self.bonds() {
            let _ = builder.add_bond(bond.atom1, bond.atom2, bond.order);
        }
        let bond_idx = builder.add_bond(a, b, order)?;
        Ok((builder.build(), bond_idx))
    }

    /// Return a new `Molecule` with the formal charge of atom `idx` changed.
    pub fn with_atom_charge(&self, idx: AtomIdx, charge: i8) -> Molecule {
        let mut builder = MoleculeBuilder::new();
        for (aidx, atom) in self.atoms() {
            let mut a = atom.clone();
            if aidx == idx { a.charge = charge; }
            builder.add_atom(a);
        }
        for (_, bond) in self.bonds() {
            let _ = builder.add_bond(bond.atom1, bond.atom2, bond.order);
        }
        builder.build()
    }

    /// Return a new `Molecule` with the element of atom `idx` changed.
    ///
    /// Chirality and hydrogen count are reset to `None` when the element
    /// changes, since those properties are element-specific.
    pub fn with_atom_element(&self, idx: AtomIdx, el: Element) -> Molecule {
        let mut builder = MoleculeBuilder::new();
        for (aidx, atom) in self.atoms() {
            let mut a = atom.clone();
            if aidx == idx {
                a.element = el;
                // Reset element-specific fields so valence stays consistent.
                a.chirality = crate::atom::Chirality::None;
                a.hydrogen_count = None;
                a.aromatic = false;
            }
            builder.add_atom(a);
        }
        for (_, bond) in self.bonds() {
            let _ = builder.add_bond(bond.atom1, bond.atom2, bond.order);
        }
        builder.build()
    }

    /// Return a new `Molecule` with atom `idx` and all bonds involving it
    /// removed.  Atom indices of survivors shift down past the removed slot.
    ///
    /// The returned tuple also includes a mapping from **old** `AtomIdx` to
    /// **new** `AtomIdx` (indices that fall below `idx` are unchanged; indices
    /// above `idx` decrease by 1).
    pub fn with_atom_removed(&self, idx: AtomIdx) -> (Molecule, Vec<Option<AtomIdx>>) {
        let n = self.atom_count();
        let removed = idx.0 as usize;

        // Build old→new index table.
        let mut remap: Vec<Option<AtomIdx>> = vec![None; n];
        let mut new_pos = 0u32;
        for old in 0..n {
            if old == removed {
                continue;
            }
            remap[old] = Some(AtomIdx(new_pos));
            new_pos += 1;
        }

        let mut builder = MoleculeBuilder::new();
        for (aidx, atom) in self.atoms() {
            if aidx == idx { continue; }
            builder.add_atom(atom.clone());
        }
        for (_, bond) in self.bonds() {
            if bond.atom1 == idx || bond.atom2 == idx { continue; }
            if let (Some(a1), Some(a2)) = (remap[bond.atom1.0 as usize], remap[bond.atom2.0 as usize]) {
                let _ = builder.add_bond(a1, a2, bond.order);
            }
        }
        (builder.build(), remap)
    }

    /// Return a new `Molecule` with bond `idx` removed.
    ///
    /// Atom indices are unchanged.  Bond indices of survivors shift down.
    pub fn with_bond_removed(&self, idx: BondIdx) -> Molecule {
        let mut builder = MoleculeBuilder::new();
        for (_, atom) in self.atoms() {
            builder.add_atom(atom.clone());
        }
        for (bidx, bond) in self.bonds() {
            if bidx == idx { continue; }
            let _ = builder.add_bond(bond.atom1, bond.atom2, bond.order);
        }
        builder.build()
    }
}

/// Builder for constructing a [`Molecule`] incrementally.
///
/// Usage: add atoms, add bonds, then call `build()`.
#[derive(Default)]
pub struct MoleculeBuilder {
    atoms: Vec<Atom>,
    bonds: Vec<BondEntry>,
    adjacency: Vec<Vec<(AtomIdx, BondIdx)>>,
}

impl MoleculeBuilder {
    pub fn new() -> Self {
        Self::default()
    }

    /// Read-only reference to an atom already added to the builder.
    ///
    /// Used by the SMILES parser to infer implicit bond types without
    /// consuming the builder (e.g. aromatic-aromatic → Aromatic bond).
    ///
    /// # Panics
    /// Panics if `idx` is out of range.
    pub fn atom_at(&self, idx: AtomIdx) -> &Atom {
        &self.atoms[idx.0 as usize]
    }

    /// Number of atoms added so far.
    pub fn atom_count(&self) -> usize {
        self.atoms.len()
    }

    /// Iterate over already-added neighbors of `idx` as `(bond_idx, neighbor_atom_idx)`.
    /// Used by kekulization tests to check whether a bond already exists in the builder.
    pub fn atom_neighbors(&self, idx: AtomIdx) -> impl Iterator<Item = (BondIdx, AtomIdx)> + '_ {
        self.adjacency[idx.0 as usize].iter().map(|&(nb, bidx)| (bidx, nb))
    }

    /// Add an atom and return its index.
    pub fn add_atom(&mut self, atom: Atom) -> AtomIdx {
        let idx = AtomIdx(self.atoms.len() as u32);
        self.atoms.push(atom);
        self.adjacency.push(Vec::new());
        idx
    }

    /// Add a bond between two existing atoms.
    ///
    /// Returns an error if either atom index is invalid or if the bond already exists.
    pub fn add_bond(&mut self, a: AtomIdx, b: AtomIdx, order: BondOrder) -> Result<BondIdx, MolError> {
        let n = self.atoms.len() as u32;
        if a.0 >= n { return Err(MolError::InvalidAtomIdx(a)); }
        if b.0 >= n { return Err(MolError::InvalidAtomIdx(b)); }

        // Check for duplicate
        for &(nb, _) in &self.adjacency[a.0 as usize] {
            if nb == b {
                return Err(MolError::DuplicateBond(a, b));
            }
        }

        let bidx = BondIdx(self.bonds.len() as u32);
        self.bonds.push(BondEntry { atom1: a, atom2: b, order });
        self.adjacency[a.0 as usize].push((b, bidx));
        self.adjacency[b.0 as usize].push((a, bidx));
        Ok(bidx)
    }

    /// Consume the builder and return an immutable [`Molecule`].
    pub fn build(self) -> Molecule {
        Molecule {
            atoms: self.atoms,
            bonds: self.bonds,
            adjacency: self.adjacency,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::atom::Atom;
    use crate::element::Element;

    fn ethane() -> Molecule {
        let mut b = MoleculeBuilder::new();
        let c1 = b.add_atom(Atom::new(Element::C));
        let c2 = b.add_atom(Atom::new(Element::C));
        b.add_bond(c1, c2, BondOrder::Single).unwrap();
        b.build()
    }

    #[test]
    fn test_basic_molecule() {
        let mol = ethane();
        assert_eq!(mol.atom_count(), 2);
        assert_eq!(mol.bond_count(), 1);
    }

    #[test]
    fn test_adjacency() {
        let mol = ethane();
        let neighbors: Vec<_> = mol.neighbors(AtomIdx(0)).collect();
        assert_eq!(neighbors.len(), 1);
        assert_eq!(neighbors[0].0, AtomIdx(1));
    }

    #[test]
    fn test_bond_between() {
        let mol = ethane();
        assert!(mol.bond_between(AtomIdx(0), AtomIdx(1)).is_some());
        assert!(mol.bond_between(AtomIdx(1), AtomIdx(0)).is_some());
    }

    #[test]
    fn test_duplicate_bond_error() {
        let mut b = MoleculeBuilder::new();
        let c1 = b.add_atom(Atom::new(Element::C));
        let c2 = b.add_atom(Atom::new(Element::C));
        b.add_bond(c1, c2, BondOrder::Single).unwrap();
        let err = b.add_bond(c1, c2, BondOrder::Double);
        assert!(matches!(err, Err(MolError::DuplicateBond(_, _))));
    }

    #[test]
    fn test_formula() {
        let mut b = MoleculeBuilder::new();
        let c = b.add_atom(Atom::new(Element::C));
        let n = b.add_atom(Atom::new(Element::N));
        b.add_bond(c, n, BondOrder::Single).unwrap();
        let mol = b.build();
        assert_eq!(mol.formula(), "CN");
    }
}