use crate::Molecule;
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub enum ExtensionValue {
Scalar(f64),
Vector(Vec<f64>),
Text(String),
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum ExtensionError {
InvalidInput(String),
Failed(String),
}
impl core::fmt::Display for ExtensionError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::InvalidInput(message) => write!(f, "invalid extension input: {message}"),
Self::Failed(message) => write!(f, "extension failed: {message}"),
}
}
}
impl std::error::Error for ExtensionError {}
pub trait MoleculeExtension: Send + Sync {
fn id(&self) -> &str;
fn version(&self) -> u32;
fn run(&self, molecule: &Molecule) -> Result<ExtensionValue, ExtensionError>;
}
#[derive(Default)]
pub struct ExtensionRegistry {
extensions: Vec<Box<dyn MoleculeExtension>>,
}
impl ExtensionRegistry {
pub fn new() -> Self {
Self::default()
}
pub fn register<E>(&mut self, extension: E) -> Result<(), ExtensionError>
where
E: MoleculeExtension + 'static,
{
if self
.extensions
.iter()
.any(|existing| existing.id() == extension.id())
{
return Err(ExtensionError::Failed(format!(
"duplicate extension id: {}",
extension.id()
)));
}
self.extensions.push(Box::new(extension));
Ok(())
}
pub fn extensions(&self) -> impl Iterator<Item = &dyn MoleculeExtension> {
self.extensions.iter().map(Box::as_ref)
}
pub fn run(&self, id: &str, molecule: &Molecule) -> Result<ExtensionValue, ExtensionError> {
self.extensions
.iter()
.find(|extension| extension.id() == id)
.ok_or_else(|| ExtensionError::Failed(format!("unknown extension id: {id}")))?
.run(molecule)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Atom, Element, MoleculeBuilder};
struct AtomCount;
impl MoleculeExtension for AtomCount {
fn id(&self) -> &str {
"test.atom-count.v1"
}
fn version(&self) -> u32 {
1
}
fn run(&self, molecule: &Molecule) -> Result<ExtensionValue, ExtensionError> {
Ok(ExtensionValue::Scalar(molecule.atom_count() as f64))
}
}
fn ethanol() -> Molecule {
let mut builder = MoleculeBuilder::new();
let c1 = builder.add_atom(Atom::new(Element::C));
let c2 = builder.add_atom(Atom::new(Element::C));
builder.add_bond(c1, c2, crate::BondOrder::Single).unwrap();
builder.build()
}
#[test]
fn registry_runs_extension_and_preserves_order() {
let mut registry = ExtensionRegistry::new();
registry.register(AtomCount).unwrap();
assert_eq!(
registry
.extensions()
.map(|extension| extension.id())
.collect::<Vec<_>>(),
["test.atom-count.v1"]
);
assert_eq!(
registry.run("test.atom-count.v1", ðanol()),
Ok(ExtensionValue::Scalar(2.0))
);
}
#[test]
fn registry_rejects_duplicate_and_unknown_ids() {
let mut registry = ExtensionRegistry::new();
registry.register(AtomCount).unwrap();
assert!(matches!(
registry.register(AtomCount),
Err(ExtensionError::Failed(message)) if message.contains("duplicate")
));
assert!(matches!(
registry.run("missing.v1", ðanol()),
Err(ExtensionError::Failed(message)) if message.contains("unknown")
));
}
}