use std::collections::BTreeMap;
use std::sync::OnceLock;
use nucleide_nuclei::{element_z, NuclideId};
use thiserror::Error;
use crate::Material;
pub type FormulaResult<T> = std::result::Result<T, FormulaError>;
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum FormulaError {
#[error("formula syntax error at byte {pos}: {message}")]
ParseError {
pos: usize,
message: String,
},
#[error("unknown element symbol `{0}`")]
UnknownElement(String),
#[error("no natural abundance data for element Z={0}")]
NoAbundanceData(u32),
#[error(transparent)]
Core(#[from] crate::Error),
}
struct Parser<'a> {
src: &'a [u8],
pos: usize,
}
impl Parser<'_> {
fn parse_error(&self, pos: usize, message: impl Into<String>) -> FormulaError {
FormulaError::ParseError {
pos,
message: message.into(),
}
}
fn error_here(&self, message: impl Into<String>) -> FormulaError {
self.parse_error(self.pos, message)
}
fn peek(&self) -> Option<u8> {
self.src.get(self.pos).copied()
}
fn at_hyphen_dot(&self) -> bool {
match self.peek() {
Some(b'.') => true,
Some(0xC2) => self.src.get(self.pos + 1) == Some(&0xB7),
_ => false,
}
}
fn advance_over_separator(&mut self) {
self.pos += if self.src[self.pos] == b'.' { 1 } else { 2 };
}
fn take_count(&mut self) -> FormulaResult<Option<f64>> {
let start = self.pos;
while self.peek().is_some_and(|c| c.is_ascii_digit()) {
self.pos += 1;
}
if start == self.pos {
return Ok(None);
}
let text = std::str::from_utf8(&self.src[start..self.pos]).unwrap_or_default();
text.parse::<u64>()
.map(|n| Some(n as f64))
.map_err(|_| self.parse_error(start, "count too large"))
}
fn take_element(&mut self) -> FormulaResult<(u32, f64)> {
let start = self.pos;
self.pos += 1;
let two_letter = self
.src
.get(start..start + 2)
.filter(|bytes| bytes[1].is_ascii_lowercase())
.and_then(|bytes| std::str::from_utf8(bytes).ok());
let one_letter = std::str::from_utf8(&self.src[start..start + 1]).ok();
let (z, matched_len) = match two_letter.and_then(element_z) {
Some(z) => (Some(z), 2),
None => (one_letter.and_then(element_z), 1),
};
let z = z.ok_or_else(|| {
let candidate = two_letter.unwrap_or_else(|| one_letter.unwrap_or("?"));
FormulaError::UnknownElement(candidate.to_string())
})?;
self.pos = start + matched_len;
let count = self.take_count()?.unwrap_or(1.0);
Ok((z, count))
}
fn take_units(&mut self, out: &mut BTreeMap<u32, f64>, scale: f64) -> FormulaResult<()> {
while let Some(c) = self.peek() {
match c {
b')' | b'.' => break,
0xC2 if self.at_hyphen_dot() => break,
b'(' => {
self.pos += 1;
let mut inner = BTreeMap::new();
self.take_units(&mut inner, 1.0)?;
if self.peek() != Some(b')') {
return Err(self.error_here("unbalanced parenthesis: expected `)`"));
}
self.pos += 1;
let mult = self.take_count()?.unwrap_or(1.0);
for (z, count) in inner {
*out.entry(z).or_insert(0.0) += count * mult * scale;
}
}
b'A'..=b'Z' => {
let (z, count) = self.take_element()?;
if count > 0.0 {
*out.entry(z).or_insert(0.0) += count * scale;
}
}
b'0'..=b'9' => {
return Err(self.error_here("count without a preceding element"));
}
_ => {
let ch = std::str::from_utf8(&self.src[self.pos..])
.unwrap_or("?")
.chars()
.next()
.unwrap_or('?');
return Err(self.error_here(format!("unexpected character `{ch}`")));
}
}
}
Ok(())
}
}
pub fn parse_formula(formula: &str) -> FormulaResult<Vec<(u32, f64)>> {
let trimmed = formula.trim();
let mut p = Parser {
src: trimmed.as_bytes(),
pos: 0,
};
let mut acc = BTreeMap::new();
let mut first = true;
while p.pos < trimmed.len() {
if !first {
if !p.at_hyphen_dot() {
if p.peek() == Some(b')') {
return Err(p.error_here("unbalanced parenthesis: unexpected `)`"));
}
return Err(p.error_here("expected `.` or `·` hydrate separator"));
}
p.advance_over_separator();
}
if p.peek().is_some_and(|c| c.is_ascii_digit()) {
if first {
return Err(p.parse_error(p.pos, "formula must not begin with a digit"));
}
let mult = p.take_count()?.unwrap_or(1.0);
p.take_units(&mut acc, mult)?;
} else {
p.take_units(&mut acc, 1.0)?;
}
first = false;
}
if first {
return Err(FormulaError::ParseError {
pos: 0,
message: "empty formula".to_string(),
});
}
Ok(acc.into_iter().collect())
}
pub trait AbundanceProvider {
fn natural_isotopes(&self, z: u32) -> Option<Vec<(NuclideId, f64)>>;
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct NoAbundances;
impl AbundanceProvider for NoAbundances {
fn natural_isotopes(&self, _z: u32) -> Option<Vec<(NuclideId, f64)>> {
None
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct NaturalAbundances;
impl NaturalAbundances {
fn groups() -> &'static BTreeMap<u32, Vec<(NuclideId, f64)>> {
static GROUPS: OnceLock<BTreeMap<u32, Vec<(NuclideId, f64)>>> = OnceLock::new();
GROUPS.get_or_init(|| {
let mut groups: BTreeMap<u32, Vec<(NuclideId, f64)>> = BTreeMap::new();
for (&nucid, &frac) in nucleide_nuclei::data::abundance_table() {
if frac > 0.0 {
let id = NuclideId::from_nucid(nucid);
groups.entry(id.z()).or_default().push((id, frac));
}
}
groups
})
}
}
impl AbundanceProvider for NaturalAbundances {
fn natural_isotopes(&self, z: u32) -> Option<Vec<(NuclideId, f64)>> {
Self::groups().get(&z).filter(|v| !v.is_empty()).cloned()
}
}
fn is_elemental(id: NuclideId) -> bool {
id.a() == 0 && id.state() == 0
}
fn ground_mass(masses: &impl crate::MassProvider, id: NuclideId) -> Option<f64> {
masses
.mass(id.nucid())
.or_else(|| masses.mass(id.nucid() - id.state()))
}
impl Material {
pub fn from_formula(
formula: &str,
masses: &impl crate::MassProvider,
abundances: &impl AbundanceProvider,
density: Option<f64>,
) -> FormulaResult<Self> {
let elements = parse_formula(formula)?;
let mut atoms = Vec::new();
for &(z, count) in &elements {
let isotopes = abundances
.natural_isotopes(z)
.ok_or(FormulaError::NoAbundanceData(z))?;
let total: f64 = isotopes.iter().map(|(_, x)| x).sum();
if total <= 0.0 {
return Err(FormulaError::NoAbundanceData(z));
}
for (id, frac) in isotopes {
atoms.push((id, count * frac / total));
}
}
Ok(Material::from_atom_frac(&atoms, masses, density)?)
}
pub fn expand_elements(
&mut self,
masses: &impl crate::MassProvider,
abundances: &impl AbundanceProvider,
) -> FormulaResult<()> {
let mut expanded = BTreeMap::new();
for (&id, &grams) in &self.comp {
if !is_elemental(id) {
expanded.insert(id, grams);
continue;
}
let z = id.z();
let isotopes = abundances
.natural_isotopes(z)
.ok_or(FormulaError::NoAbundanceData(z))?;
let total: f64 = isotopes.iter().map(|(_, x)| x).sum();
if total <= 0.0 {
return Err(FormulaError::NoAbundanceData(z));
}
let mean_mass = isotopes
.iter()
.map(|&(iso, x)| {
ground_mass(masses, iso)
.ok_or(crate::Error::MissingMass(iso))
.map(|m| x / total * m)
})
.sum::<crate::Result<f64>>()
.map_err(FormulaError::from)?;
for (iso, x) in isotopes {
let m = ground_mass(masses, iso).expect("checked by mean_mass loop above");
expanded.insert(iso, grams * (x / total) * m / mean_mass);
}
}
self.comp = expanded;
Ok(())
}
pub fn collapse_elements(&self) -> Self {
let mut comp = BTreeMap::new();
for (&id, &grams) in &self.comp {
let key = if is_elemental(id) {
id
} else {
NuclideId::from_nucid(id.z() * 10_000_000)
};
*comp.entry(key).or_insert(0.0) += grams;
}
let mut out = Material::new();
out.comp = comp;
out.set_density(self.density());
out.set_metadata(self.metadata().cloned());
out
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Ame2020;
fn counts(formula: &str) -> Vec<(u32, f64)> {
parse_formula(formula).unwrap()
}
fn assert_counts(formula: &str, expected: &[(u32, f64)]) {
assert_eq!(counts(formula), expected.to_vec(), "for `{formula}`");
}
#[test]
fn parses_water_and_glucose() {
assert_counts("H2O", &[(1, 2.0), (8, 1.0)]);
assert_counts("C6H12O6", &[(1, 12.0), (6, 6.0), (8, 6.0)]);
}
#[test]
fn parses_grouped_and_nested_formulas() {
assert_counts("Ca(OH)2", &[(1, 2.0), (8, 2.0), (20, 1.0)]);
assert_counts("Fe2(SO4)3", &[(8, 12.0), (16, 3.0), (26, 2.0)]);
assert_counts("Mg(NO2)2", &[(7, 2.0), (8, 4.0), (12, 1.0)]);
assert_counts("U((C)3)2", &[(6, 6.0), (92, 1.0)]);
}
#[test]
fn parses_chained_groups_as_single_elements() {
assert_counts("CH3(CH2)6CH3", &[(1, 18.0), (6, 8.0)]);
}
#[test]
fn parses_multi_digit_counts_and_hydrates() {
assert_counts("C12H22O11", &[(1, 22.0), (6, 12.0), (8, 11.0)]);
let dot = counts("CuSO4·5H2O");
let ascii = counts("CuSO4.5H2O");
assert_eq!(dot, ascii);
assert_eq!(dot, vec![(1, 10.0), (8, 9.0), (16, 1.0), (29, 1.0)]);
assert_counts("H2O.2H2O", &[(1, 6.0), (8, 3.0)]);
}
#[test]
fn tolerates_surrounding_whitespace_only() {
assert_eq!(counts(" H2O "), counts("H2O"));
}
#[test]
fn rejects_unbalanced_parentheses() {
let err = parse_formula("(H2O").unwrap_err();
assert!(
matches!(err, FormulaError::ParseError { pos: 4, .. }),
"{err}"
);
let err = parse_formula("H2O)").unwrap_err();
assert!(matches!(err, FormulaError::ParseError { .. }), "{err}");
}
#[test]
fn rejects_unknown_symbol() {
match parse_formula("Xx2O").unwrap_err() {
FormulaError::UnknownElement(s) => assert_eq!(s, "Xx"),
other => panic!("{other:?}"),
}
match parse_formula("Q").unwrap_err() {
FormulaError::UnknownElement(s) => assert_eq!(s, "Q"),
other => panic!("{other:?}"),
}
}
#[test]
fn rejects_leading_digit_empty_and_stray_chars() {
let err = parse_formula("2H2O").unwrap_err();
assert!(
matches!(err, FormulaError::ParseError { pos: 0, .. }),
"{err}"
);
assert!(matches!(
parse_formula("").unwrap_err(),
FormulaError::ParseError { .. }
));
let err = parse_formula("H2 O").unwrap_err();
assert!(
matches!(err, FormulaError::ParseError { pos: 2, .. }),
"{err}"
);
let err = parse_formula("H1O-1").unwrap_err();
assert!(matches!(err, FormulaError::ParseError { .. }));
}
#[test]
fn from_formula_water_has_natural_isotopes_and_two_thirds_hydrogen() {
let water = Material::from_formula("H2O", &Ame2020, &NaturalAbundances, Some(1.0)).unwrap();
assert!(water
.comp
.contains_key(&NuclideId::from_name("H1").unwrap()));
for o in ["O16", "O17", "O18"] {
assert!(
water.comp.contains_key(&NuclideId::from_name(o).unwrap()),
"missing {o}"
);
}
let af = water.atom_fractions(&Ame2020).unwrap();
let h: f64 = af
.iter()
.filter(|(id, _)| id.z() == 1)
.map(|(_, f)| f)
.sum();
assert!((h - 2.0 / 3.0).abs() < 1e-12, "{h}");
let o: f64 = af
.iter()
.filter(|(id, _)| id.z() == 8)
.map(|(_, f)| f)
.sum();
assert!((o - 1.0 / 3.0).abs() < 1e-12);
assert_eq!(water.density(), Some(1.0));
}
#[test]
fn from_formula_h2so4() {
let mat = Material::from_formula("H2SO4", &Ame2020, &NaturalAbundances, None).unwrap();
let af = mat.atom_fractions(&Ame2020).unwrap();
assert!(!af.is_empty(), "H2SO4 atom fractions should not be empty");
for nuc in ["H1", "H2", "O16", "O17", "O18", "S32", "S33", "S34", "S36"] {
assert!(
af.contains_key(&NuclideId::from_name(nuc).unwrap()),
"missing {nuc}"
);
}
}
#[test]
fn from_formula_rejects_bad_input_and_missing_data() {
match Material::from_formula("Xx", &Ame2020, &NaturalAbundances, None).unwrap_err() {
FormulaError::UnknownElement(s) => assert_eq!(s, "Xx"),
other => panic!("{other:?}"),
}
assert!(matches!(
Material::from_formula("U", &Ame2020, &NoAbundances, None).unwrap_err(),
FormulaError::NoAbundanceData(92)
));
}
#[test]
fn expand_then_collapse_round_trips_an_elemental_material() {
let mut mat = Material::new();
mat.add_nuclide(NuclideId::from_nucid(10_000_000), 2.0); mat.add_nuclide(NuclideId::from_nucid(80_000_000), 16.0); let original = mat.clone();
mat.expand_elements(&Ame2020, &NaturalAbundances).unwrap();
assert!(!mat.comp.contains_key(&NuclideId::from_nucid(80_000_000)));
assert!(mat.comp.contains_key(&NuclideId::from_name("H1").unwrap()));
assert!(mat.comp.contains_key(&NuclideId::from_name("H2").unwrap()));
assert!(mat.comp.contains_key(&NuclideId::from_name("O18").unwrap()));
let back = mat.collapse_elements();
assert_eq!(
back.comp.keys().copied().collect::<Vec<_>>(),
original.comp.keys().copied().collect::<Vec<_>>()
);
for (id, m0) in &original.comp {
let m1 = back.comp[id];
assert!((m0 - m1).abs() < 1e-9 * m0.abs(), "{id}: {m0} vs {m1}");
}
}
#[test]
fn expand_preserves_entry_masses_and_leaves_named_nuclides_alone() {
let mut mat = Material::new();
mat.add_nuclide(NuclideId::from_nucid(10_000_000), 18.0); mat.add_nuclide(NuclideId::from_name("Fe56").unwrap(), 5.0);
mat.expand_elements(&Ame2020, &NaturalAbundances).unwrap();
close(mat.mass(), 23.0);
close(
mat.remove_nuclide(NuclideId::from_name("Fe56").unwrap())
.unwrap(),
5.0,
);
let h: f64 = mat.comp.values().sum();
close(h, 18.0);
}
#[test]
fn expand_without_abundances_errors_with_z() {
let mut mat = Material::new();
mat.add_nuclide(NuclideId::from_nucid(920_000_000), 1.0);
match mat.expand_elements(&Ame2020, &NoAbundances).unwrap_err() {
FormulaError::NoAbundanceData(z) => assert_eq!(z, 92),
other => panic!("{other:?}"),
}
}
#[test]
fn collapse_folds_named_nuclides_into_placeholder_keys() {
let mut mat = Material::new();
mat.add_nuclide(NuclideId::from_name("U235").unwrap(), 3.0);
mat.add_nuclide(NuclideId::from_name("U238").unwrap(), 1.0);
mat.set_density(Some(19.1));
let collapsed = mat.collapse_elements();
let key = NuclideId::from_nucid(920_000_000);
assert_eq!(collapsed.comp.len(), 1);
close(collapsed.comp[&key], 4.0);
assert_eq!(
key.nucid(),
nucleide_nuclei::element_z("U").unwrap() * 10_000_000
);
assert_eq!(collapsed.density(), Some(19.1));
}
fn close(a: f64, b: f64) {
assert!((a - b).abs() < 1e-12, "{a} != {b}");
}
}