use std::collections::HashMap;
use chematic_core::{AtomIdx, BondOrder, Molecule, implicit_hcount};
use chematic_perception::RingSet;
use chematic_perception::find_sssr;
use crate::query::{AtomPrimitive, AtomQuery, BondPrimitive, BondQuery, QueryMolecule};
struct EvalCtx<'a> {
mol: &'a Molecule,
rings: RingSet,
config: &'a MatchConfig,
visit_budget: std::cell::Cell<u64>,
}
impl<'a> EvalCtx<'a> {
fn new(mol: &'a Molecule, config: &'a MatchConfig) -> Self {
Self {
mol,
rings: find_sssr(mol),
config,
visit_budget: std::cell::Cell::new(
config.max_visit_budget.unwrap_or(u64::MAX),
),
}
}
}
#[derive(Debug, Clone)]
pub struct MatchConfig {
pub max_matches: Option<usize>,
pub use_chirality: bool,
pub use_isotopes: bool,
pub uniquify: bool,
pub max_visit_budget: Option<u64>,
}
impl Default for MatchConfig {
fn default() -> Self {
Self {
max_matches: None,
use_chirality: false,
use_isotopes: false,
uniquify: true,
max_visit_budget: None,
}
}
}
pub fn find_matches(query: &QueryMolecule, mol: &Molecule) -> Vec<HashMap<usize, AtomIdx>> {
find_matches_with_config(query, mol, &MatchConfig::default())
}
pub fn find_matches_with_config(
query: &QueryMolecule,
mol: &Molecule,
config: &MatchConfig,
) -> Vec<HashMap<usize, AtomIdx>> {
if query.atoms.is_empty() {
return vec![];
}
let ctx = EvalCtx::new(mol, config);
let mut mapping: HashMap<usize, AtomIdx> = HashMap::new();
let mut results: Vec<HashMap<usize, AtomIdx>> = Vec::new();
match_recursive(query, &ctx, &mut mapping, &mut results, config.max_matches);
if config.uniquify {
let mut seen = std::collections::HashSet::new();
results.retain(|m| {
let mut key: Vec<u32> = m.values().map(|idx| idx.0).collect();
key.sort_unstable();
seen.insert(key)
});
}
results
}
fn next_unmapped(mapping: &HashMap<usize, AtomIdx>, query_len: usize) -> usize {
(0..query_len)
.find(|i| !mapping.contains_key(i))
.unwrap() }
fn match_recursive(
query: &QueryMolecule,
ctx: &EvalCtx<'_>,
mapping: &mut HashMap<usize, AtomIdx>,
results: &mut Vec<HashMap<usize, AtomIdx>>,
max: Option<usize>,
) {
if max.is_some_and(|m| results.len() >= m) {
return;
}
let remaining = ctx.visit_budget.get();
if remaining == 0 {
return;
}
ctx.visit_budget.set(remaining - 1);
if mapping.len() == query.atoms.len() {
results.push(mapping.clone());
return;
}
let q_next = next_unmapped(mapping, query.atoms.len());
let used_targets: std::collections::HashSet<AtomIdx> = mapping.values().copied().collect();
for t in 0..ctx.mol.atom_count() {
if max.is_some_and(|m| results.len() >= m) {
break;
}
let t_idx = AtomIdx(t as u32);
if used_targets.contains(&t_idx) {
continue;
}
if !eval_atom_query(&query.atoms[q_next].query, t_idx, ctx) {
continue;
}
if !bonds_compatible(q_next, t_idx, mapping, query, ctx) {
continue;
}
mapping.insert(q_next, t_idx);
match_recursive(query, ctx, mapping, results, max);
mapping.remove(&q_next);
}
}
fn bonds_compatible(
q: usize,
t: AtomIdx,
mapping: &HashMap<usize, AtomIdx>,
query: &QueryMolecule,
ctx: &EvalCtx<'_>,
) -> bool {
for &(bond_idx, q_nb) in &query.adj[q] {
if let Some(&t_nb) = mapping.get(&q_nb) {
match ctx.mol.bond_between(t, t_nb) {
None => return false,
Some((_bidx, bond_entry)) => {
let bq = &query.bonds[bond_idx].query;
if !eval_bond_query(bq, bond_entry.order, t, t_nb, ctx) {
return false;
}
}
}
}
}
true
}
fn eval_atom_query(q: &AtomQuery, idx: AtomIdx, ctx: &EvalCtx<'_>) -> bool {
match q {
AtomQuery::Primitive(p) => eval_atom_primitive(p, idx, ctx),
AtomQuery::And(a, b) => eval_atom_query(a, idx, ctx) && eval_atom_query(b, idx, ctx),
AtomQuery::Or(a, b) => eval_atom_query(a, idx, ctx) || eval_atom_query(b, idx, ctx),
AtomQuery::Not(a) => !eval_atom_query(a, idx, ctx),
}
}
fn eval_atom_primitive(p: &AtomPrimitive, idx: AtomIdx, ctx: &EvalCtx<'_>) -> bool {
let atom = ctx.mol.atom(idx);
match p {
AtomPrimitive::AtomicNum(n) => atom.element.atomic_number() == *n,
AtomPrimitive::Symbol(s) => atom.element.symbol() == s.as_str(),
AtomPrimitive::Aromatic(a) => atom.aromatic == *a,
AtomPrimitive::Charge(c) => atom.charge == *c,
AtomPrimitive::HCount(h) => eval_hcount(idx, ctx, *h),
AtomPrimitive::ImplicitHCount(h) => implicit_hcount(ctx.mol, idx) == *h,
AtomPrimitive::Degree(d) => ctx.mol.neighbors(idx).count() as u8 == *d,
AtomPrimitive::RingMembership(r) => ctx.rings.contains_atom(idx) == *r,
AtomPrimitive::RingSize(n) => ctx
.rings
.rings()
.iter()
.any(|ring| ring.len() == *n as usize && ring.contains(&idx)),
AtomPrimitive::Wildcard => true,
AtomPrimitive::Recursive(sub_query) => has_match_anchored(sub_query, idx, ctx),
AtomPrimitive::Valence(v) => eval_valence(idx, ctx, *v),
AtomPrimitive::RingBondCount(x) => eval_ring_bond_count(idx, ctx, *x),
AtomPrimitive::TotalConnectivity(x) => {
ctx.mol.neighbors(idx).count() as u8 + implicit_hcount(ctx.mol, idx) == *x
}
AtomPrimitive::RingCount(n) => {
ctx.rings.rings().iter().filter(|r| r.contains(&idx)).count() as u8 == *n
}
AtomPrimitive::Hybridization(h) => eval_hybridization(idx, ctx, *h),
AtomPrimitive::Isotope(mass) => {
!ctx.config.use_isotopes || ctx.mol.atom(idx).isotope == Some(*mass)
}
AtomPrimitive::Chirality(kind) => eval_chirality(idx, ctx, *kind),
}
}
fn eval_hcount(idx: AtomIdx, ctx: &EvalCtx<'_>, h: u8) -> bool {
let explicit_h = ctx
.mol
.neighbors(idx)
.filter(|(nb, _)| ctx.mol.atom(*nb).element.atomic_number() == 1)
.count() as u8;
explicit_h + implicit_hcount(ctx.mol, idx) == h
}
fn eval_valence(idx: AtomIdx, ctx: &EvalCtx<'_>, v: u8) -> bool {
let bond_sum: u8 = ctx
.mol
.neighbors(idx)
.map(|(_, bid)| bond_order_int(ctx.mol.bond(bid).order))
.sum();
bond_sum + implicit_hcount(ctx.mol, idx) == v
}
fn eval_ring_bond_count(idx: AtomIdx, ctx: &EvalCtx<'_>, x: u8) -> bool {
let count = ctx
.mol
.neighbors(idx)
.filter(|(nb, _)| {
ctx.rings
.rings()
.iter()
.any(|ring| ring.contains(&idx) && ring.contains(nb))
})
.count() as u8;
count == x
}
fn eval_hybridization(idx: AtomIdx, ctx: &EvalCtx<'_>, h: u8) -> bool {
let atom = ctx.mol.atom(idx);
let hyb = if atom.aromatic {
2u8
} else {
let mut has_triple = false;
let mut has_double = false;
for (_, bid) in ctx.mol.neighbors(idx) {
match ctx.mol.bond(bid).order {
BondOrder::Triple => { has_triple = true; break; }
BondOrder::Double => has_double = true,
_ => {}
}
}
if has_triple { 1 } else if has_double { 2 } else { 3 }
};
hyb == h
}
fn eval_chirality(idx: AtomIdx, ctx: &EvalCtx<'_>, kind: u8) -> bool {
if !ctx.config.use_chirality {
return true;
}
use chematic_core::Chirality;
let c = ctx.mol.atom(idx).chirality;
match kind {
1 => c == Chirality::CounterClockwise,
2 => c == Chirality::Clockwise,
_ => c != Chirality::None,
}
}
fn has_match_anchored(query: &QueryMolecule, anchor: AtomIdx, ctx: &EvalCtx<'_>) -> bool {
if query.atoms.is_empty() {
return false;
}
if !eval_atom_query(&query.atoms[0].query, anchor, ctx) {
return false;
}
let mut mapping = HashMap::new();
mapping.insert(0usize, anchor);
if query.atoms.len() == 1 {
return true;
}
has_match_recursive(query, ctx, &mut mapping)
}
fn has_match_recursive(
query: &QueryMolecule,
ctx: &EvalCtx<'_>,
mapping: &mut HashMap<usize, AtomIdx>,
) -> bool {
let remaining = ctx.visit_budget.get();
if remaining == 0 {
return false;
}
ctx.visit_budget.set(remaining - 1);
if mapping.len() == query.atoms.len() {
return true;
}
let q_next = next_unmapped(mapping, query.atoms.len());
let used_targets: std::collections::HashSet<AtomIdx> = mapping.values().copied().collect();
for t in 0..ctx.mol.atom_count() {
let t_idx = AtomIdx(t as u32);
if used_targets.contains(&t_idx) {
continue;
}
if !eval_atom_query(&query.atoms[q_next].query, t_idx, ctx) {
continue;
}
if !bonds_compatible(q_next, t_idx, mapping, query, ctx) {
continue;
}
mapping.insert(q_next, t_idx);
if has_match_recursive(query, ctx, mapping) {
mapping.remove(&q_next);
return true;
}
mapping.remove(&q_next);
}
false
}
fn eval_bond_query(
q: &BondQuery,
order: BondOrder,
a: AtomIdx,
b: AtomIdx,
ctx: &EvalCtx<'_>,
) -> bool {
match q {
BondQuery::Primitive(p) => eval_bond_primitive(p, order, a, b, ctx),
BondQuery::And(x, y) => {
eval_bond_query(x, order, a, b, ctx) && eval_bond_query(y, order, a, b, ctx)
}
BondQuery::Or(x, y) => {
eval_bond_query(x, order, a, b, ctx) || eval_bond_query(y, order, a, b, ctx)
}
BondQuery::Not(x) => !eval_bond_query(x, order, a, b, ctx),
BondQuery::Any => true,
}
}
fn bond_order_int(order: BondOrder) -> u8 {
match order {
BondOrder::Zero => 0,
BondOrder::Single
| BondOrder::Up
| BondOrder::Down
| BondOrder::Aromatic
| BondOrder::Dative
| BondOrder::QueryAny
| BondOrder::QuerySingleOrDouble
| BondOrder::QuerySingleOrAromatic
| BondOrder::QueryDoubleOrAromatic => 1,
BondOrder::Double => 2,
BondOrder::Triple => 3,
BondOrder::Quadruple => 4,
}
}
fn eval_bond_primitive(
p: &BondPrimitive,
order: BondOrder,
a: AtomIdx,
b: AtomIdx,
ctx: &EvalCtx<'_>,
) -> bool {
match p {
BondPrimitive::Single => {
matches!(
order,
BondOrder::Single
| BondOrder::Up
| BondOrder::Down
| BondOrder::QuerySingleOrDouble
| BondOrder::QuerySingleOrAromatic
)
}
BondPrimitive::Double => matches!(
order,
BondOrder::Double | BondOrder::QuerySingleOrDouble | BondOrder::QueryDoubleOrAromatic
),
BondPrimitive::Triple => matches!(order, BondOrder::Triple),
BondPrimitive::Aromatic => matches!(
order,
BondOrder::Aromatic
| BondOrder::QuerySingleOrAromatic
| BondOrder::QueryDoubleOrAromatic
),
BondPrimitive::Any => true,
BondPrimitive::Ring => {
ctx.rings
.rings()
.iter()
.any(|ring| ring.contains(&a) && ring.contains(&b))
}
BondPrimitive::Up => matches!(order, BondOrder::Up),
BondPrimitive::Down => matches!(order, BondOrder::Down),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{find_matches, find_matches_with_config, parse_smarts};
use chematic_smiles::parse;
#[test]
fn test_isotope_ignored_by_default() {
let mol = parse("CC").unwrap();
let query = parse_smarts("[13C]").unwrap();
let matches = find_matches(&query, &mol);
assert_eq!(
matches.len(),
2,
"[13C] with use_isotopes=false should match all carbons"
);
}
#[test]
fn test_isotope_enforced_when_enabled() {
use chematic_core::{Atom, AtomIdx, BondOrder, Element, MoleculeBuilder};
let mut b = MoleculeBuilder::new();
let mut c13 = Atom::new(Element::C);
c13.isotope = Some(13);
let c13_idx = b.add_atom(c13);
let c12_idx = b.add_atom(Atom::new(Element::C));
b.add_bond(c13_idx, c12_idx, BondOrder::Single).unwrap();
let mol = b.build();
let query = parse_smarts("[13C]").unwrap();
let config = MatchConfig {
use_isotopes: true,
..MatchConfig::default()
};
let matches = find_matches_with_config(&query, &mol, &config);
assert_eq!(
matches.len(),
1,
"[13C] with use_isotopes=true should match only the 13C atom"
);
assert_eq!(matches[0][&0], AtomIdx(0));
}
#[test]
fn test_no_isotope_match_on_unlabeled() {
let mol = parse("CC").unwrap(); let query = parse_smarts("[13C]").unwrap();
let config = MatchConfig {
use_isotopes: true,
..MatchConfig::default()
};
let matches = find_matches_with_config(&query, &mol, &config);
assert_eq!(
matches.len(),
0,
"[13C] with use_isotopes=true should not match unlabeled C"
);
}
#[test]
fn test_chirality_ignored_by_default() {
let mol = parse("N[C@@H](C)C(=O)O").unwrap(); let query = parse_smarts("[C@@H]").unwrap();
let matches = find_matches(&query, &mol);
assert!(
!matches.is_empty(),
"chirality should be ignored by default"
);
}
#[test]
fn test_chirality_enforced_when_enabled() {
let mol = parse("N[C@@H](C)C(=O)O").unwrap();
let q_ccw = parse_smarts("[C@@H]").unwrap(); let q_cw = parse_smarts("[C@H]").unwrap();
let config = MatchConfig {
use_chirality: true,
..MatchConfig::default()
};
let m_ccw = find_matches_with_config(&q_ccw, &mol, &config);
let m_cw = find_matches_with_config(&q_cw, &mol, &config);
assert!(!m_ccw.is_empty(), "[C@@H] should match L-alanine (@@)");
assert!(m_cw.is_empty(), "[C@H] should not match L-alanine (@@)");
}
#[test]
fn test_chirality_d_alanine_positive_match() {
let mol = parse("N[C@H](C)C(=O)O").unwrap(); let q_cw = parse_smarts("[C@H]").unwrap(); let config = MatchConfig {
use_chirality: true,
..MatchConfig::default()
};
let m = find_matches_with_config(&q_cw, &mol, &config);
assert!(!m.is_empty(), "[C@H] should positively match D-alanine (@)");
}
#[test]
fn test_visit_budget_unlimited_default() {
let mol = parse("c1ccccc1").unwrap(); let q = parse_smarts("c1ccccc1").unwrap();
let m = find_matches(&q, &mol);
assert!(!m.is_empty(), "benzene should match aromatic ring query");
}
#[test]
fn test_visit_budget_generous_limit_finds_match() {
let mol = parse("CCO").unwrap();
let q = parse_smarts("O").unwrap();
let config = MatchConfig {
max_visit_budget: Some(10_000),
..MatchConfig::default()
};
let m = find_matches_with_config(&q, &mol, &config);
assert!(!m.is_empty(), "CCO contains O — should match within budget");
}
#[test]
fn test_visit_budget_zero_returns_empty() {
let mol = parse("CCO").unwrap();
let q = parse_smarts("O").unwrap();
let config = MatchConfig {
max_visit_budget: Some(0),
..MatchConfig::default()
};
let m = find_matches_with_config(&q, &mol, &config);
let _ = m;
}
}