1use std::collections::BTreeMap;
30use std::sync::OnceLock;
31
32use nucleide_nuclei::{element_z, NuclideId};
33use thiserror::Error;
34
35use crate::Material;
36pub type FormulaResult<T> = std::result::Result<T, FormulaError>;
38
39#[derive(Debug, Error)]
45#[non_exhaustive]
46pub enum FormulaError {
47 #[error("formula syntax error at byte {pos}: {message}")]
50 ParseError {
51 pos: usize,
53 message: String,
55 },
56 #[error("unknown element symbol `{0}`")]
58 UnknownElement(String),
59 #[error("no natural abundance data for element Z={0}")]
61 NoAbundanceData(u32),
62 #[error(transparent)]
64 Core(#[from] crate::Error),
65}
66
67struct Parser<'a> {
81 src: &'a [u8],
82 pos: usize,
83}
84
85impl Parser<'_> {
86 fn parse_error(&self, pos: usize, message: impl Into<String>) -> FormulaError {
87 FormulaError::ParseError {
88 pos,
89 message: message.into(),
90 }
91 }
92
93 fn error_here(&self, message: impl Into<String>) -> FormulaError {
94 self.parse_error(self.pos, message)
95 }
96
97 fn peek(&self) -> Option<u8> {
98 self.src.get(self.pos).copied()
99 }
100
101 fn at_hyphen_dot(&self) -> bool {
104 match self.peek() {
105 Some(b'.') => true,
106 Some(0xC2) => self.src.get(self.pos + 1) == Some(&0xB7),
107 _ => false,
108 }
109 }
110
111 fn advance_over_separator(&mut self) {
112 self.pos += if self.src[self.pos] == b'.' { 1 } else { 2 };
113 }
114
115 fn take_count(&mut self) -> FormulaResult<Option<f64>> {
117 let start = self.pos;
118 while self.peek().is_some_and(|c| c.is_ascii_digit()) {
119 self.pos += 1;
120 }
121 if start == self.pos {
122 return Ok(None);
123 }
124 let text = std::str::from_utf8(&self.src[start..self.pos]).unwrap_or_default();
125 text.parse::<u64>()
126 .map(|n| Some(n as f64))
127 .map_err(|_| self.parse_error(start, "count too large"))
128 }
129
130 fn take_element(&mut self) -> FormulaResult<(u32, f64)> {
133 let start = self.pos;
134 self.pos += 1;
135 let two_letter = self
137 .src
138 .get(start..start + 2)
139 .filter(|bytes| bytes[1].is_ascii_lowercase())
140 .and_then(|bytes| std::str::from_utf8(bytes).ok());
141 let one_letter = std::str::from_utf8(&self.src[start..start + 1]).ok();
142 let (z, matched_len) = match two_letter.and_then(element_z) {
143 Some(z) => (Some(z), 2),
144 None => (one_letter.and_then(element_z), 1),
145 };
146 let z = z.ok_or_else(|| {
147 let candidate = two_letter.unwrap_or_else(|| one_letter.unwrap_or("?"));
148 FormulaError::UnknownElement(candidate.to_string())
149 })?;
150 self.pos = start + matched_len;
151 let count = self.take_count()?.unwrap_or(1.0);
152 Ok((z, count))
153 }
154
155 fn take_units(&mut self, out: &mut BTreeMap<u32, f64>, scale: f64) -> FormulaResult<()> {
158 while let Some(c) = self.peek() {
159 match c {
160 b')' | b'.' => break,
161 0xC2 if self.at_hyphen_dot() => break,
162 b'(' => {
163 self.pos += 1;
164 let mut inner = BTreeMap::new();
165 self.take_units(&mut inner, 1.0)?;
166 if self.peek() != Some(b')') {
167 return Err(self.error_here("unbalanced parenthesis: expected `)`"));
168 }
169 self.pos += 1;
170 let mult = self.take_count()?.unwrap_or(1.0);
171 for (z, count) in inner {
172 *out.entry(z).or_insert(0.0) += count * mult * scale;
173 }
174 }
175 b'A'..=b'Z' => {
176 let (z, count) = self.take_element()?;
177 if count > 0.0 {
178 *out.entry(z).or_insert(0.0) += count * scale;
179 }
180 }
181 b'0'..=b'9' => {
182 return Err(self.error_here("count without a preceding element"));
183 }
184 _ => {
185 let ch = std::str::from_utf8(&self.src[self.pos..])
186 .unwrap_or("?")
187 .chars()
188 .next()
189 .unwrap_or('?');
190 return Err(self.error_here(format!("unexpected character `{ch}`")));
191 }
192 }
193 }
194 Ok(())
195 }
196}
197
198pub fn parse_formula(formula: &str) -> FormulaResult<Vec<(u32, f64)>> {
211 let trimmed = formula.trim();
212 let mut p = Parser {
213 src: trimmed.as_bytes(),
214 pos: 0,
215 };
216 let mut acc = BTreeMap::new();
217 let mut first = true;
218 while p.pos < trimmed.len() {
219 if !first {
220 if !p.at_hyphen_dot() {
221 if p.peek() == Some(b')') {
222 return Err(p.error_here("unbalanced parenthesis: unexpected `)`"));
223 }
224 return Err(p.error_here("expected `.` or `·` hydrate separator"));
225 }
226 p.advance_over_separator();
227 }
228 if p.peek().is_some_and(|c| c.is_ascii_digit()) {
229 if first {
230 return Err(p.parse_error(p.pos, "formula must not begin with a digit"));
231 }
232 let mult = p.take_count()?.unwrap_or(1.0);
233 p.take_units(&mut acc, mult)?;
234 } else {
235 p.take_units(&mut acc, 1.0)?;
236 }
237 first = false;
238 }
239 if first {
240 return Err(FormulaError::ParseError {
241 pos: 0,
242 message: "empty formula".to_string(),
243 });
244 }
245 Ok(acc.into_iter().collect())
246}
247
248pub trait AbundanceProvider {
257 fn natural_isotopes(&self, z: u32) -> Option<Vec<(NuclideId, f64)>>;
260}
261
262#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
267pub struct NoAbundances;
268
269impl AbundanceProvider for NoAbundances {
270 fn natural_isotopes(&self, _z: u32) -> Option<Vec<(NuclideId, f64)>> {
271 None
272 }
273}
274
275#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
278pub struct NaturalAbundances;
279
280impl NaturalAbundances {
281 fn groups() -> &'static BTreeMap<u32, Vec<(NuclideId, f64)>> {
283 static GROUPS: OnceLock<BTreeMap<u32, Vec<(NuclideId, f64)>>> = OnceLock::new();
284 GROUPS.get_or_init(|| {
285 let mut groups: BTreeMap<u32, Vec<(NuclideId, f64)>> = BTreeMap::new();
286 for (&nucid, &frac) in nucleide_nuclei::data::abundance_table() {
287 if frac > 0.0 {
288 let id = NuclideId::from_nucid(nucid);
289 groups.entry(id.z()).or_default().push((id, frac));
290 }
291 }
292 groups
293 })
294 }
295}
296
297impl AbundanceProvider for NaturalAbundances {
298 fn natural_isotopes(&self, z: u32) -> Option<Vec<(NuclideId, f64)>> {
299 Self::groups().get(&z).filter(|v| !v.is_empty()).cloned()
300 }
301}
302
303fn is_elemental(id: NuclideId) -> bool {
311 id.a() == 0 && id.state() == 0
312}
313
314fn ground_mass(masses: &impl crate::MassProvider, id: NuclideId) -> Option<f64> {
318 masses
319 .mass(id.nucid())
320 .or_else(|| masses.mass(id.nucid() - id.state()))
321}
322
323impl Material {
324 pub fn from_formula(
334 formula: &str,
335 masses: &impl crate::MassProvider,
336 abundances: &impl AbundanceProvider,
337 density: Option<f64>,
338 ) -> FormulaResult<Self> {
339 let elements = parse_formula(formula)?;
340 let mut atoms = Vec::new();
341 for &(z, count) in &elements {
342 let isotopes = abundances
343 .natural_isotopes(z)
344 .ok_or(FormulaError::NoAbundanceData(z))?;
345 let total: f64 = isotopes.iter().map(|(_, x)| x).sum();
346 if total <= 0.0 {
347 return Err(FormulaError::NoAbundanceData(z));
348 }
349 for (id, frac) in isotopes {
350 atoms.push((id, count * frac / total));
351 }
352 }
353 Ok(Material::from_atom_frac(&atoms, masses, density)?)
354 }
355
356 pub fn expand_elements(
373 &mut self,
374 masses: &impl crate::MassProvider,
375 abundances: &impl AbundanceProvider,
376 ) -> FormulaResult<()> {
377 let mut expanded = BTreeMap::new();
378 for (&id, &grams) in &self.comp {
379 if !is_elemental(id) {
380 expanded.insert(id, grams);
381 continue;
382 }
383 let z = id.z();
384 let isotopes = abundances
385 .natural_isotopes(z)
386 .ok_or(FormulaError::NoAbundanceData(z))?;
387 let total: f64 = isotopes.iter().map(|(_, x)| x).sum();
388 if total <= 0.0 {
389 return Err(FormulaError::NoAbundanceData(z));
390 }
391 let mean_mass = isotopes
393 .iter()
394 .map(|&(iso, x)| {
395 ground_mass(masses, iso)
396 .ok_or(crate::Error::MissingMass(iso))
397 .map(|m| x / total * m)
398 })
399 .sum::<crate::Result<f64>>()
400 .map_err(FormulaError::from)?;
401 for (iso, x) in isotopes {
402 let m = ground_mass(masses, iso).expect("checked by mean_mass loop above");
403 expanded.insert(iso, grams * (x / total) * m / mean_mass);
404 }
405 }
406 self.comp = expanded;
407 Ok(())
408 }
409
410 pub fn collapse_elements(&self) -> Self {
416 let mut comp = BTreeMap::new();
417 for (&id, &grams) in &self.comp {
418 let key = if is_elemental(id) {
419 id
420 } else {
421 NuclideId::from_nucid(id.z() * 10_000_000)
422 };
423 *comp.entry(key).or_insert(0.0) += grams;
424 }
425 let mut out = Material::new();
426 out.comp = comp;
427 out.set_density(self.density());
428 out.set_metadata(self.metadata().cloned());
429 out
430 }
431}
432
433#[cfg(test)]
434mod tests {
435 use super::*;
436 use crate::Ame2020;
437
438 fn counts(formula: &str) -> Vec<(u32, f64)> {
439 parse_formula(formula).unwrap()
440 }
441
442 fn assert_counts(formula: &str, expected: &[(u32, f64)]) {
443 assert_eq!(counts(formula), expected.to_vec(), "for `{formula}`");
444 }
445
446 #[test]
447 fn parses_water_and_glucose() {
448 assert_counts("H2O", &[(1, 2.0), (8, 1.0)]);
449 assert_counts("C6H12O6", &[(1, 12.0), (6, 6.0), (8, 6.0)]);
450 }
451
452 #[test]
453 fn parses_grouped_and_nested_formulas() {
454 assert_counts("Ca(OH)2", &[(1, 2.0), (8, 2.0), (20, 1.0)]);
456 assert_counts("Fe2(SO4)3", &[(8, 12.0), (16, 3.0), (26, 2.0)]);
458 assert_counts("Mg(NO2)2", &[(7, 2.0), (8, 4.0), (12, 1.0)]);
460 assert_counts("U((C)3)2", &[(6, 6.0), (92, 1.0)]);
461 }
462
463 #[test]
464 fn parses_chained_groups_as_single_elements() {
465 assert_counts("CH3(CH2)6CH3", &[(1, 18.0), (6, 8.0)]);
467 }
468
469 #[test]
470 fn parses_multi_digit_counts_and_hydrates() {
471 assert_counts("C12H22O11", &[(1, 22.0), (6, 12.0), (8, 11.0)]);
472 let dot = counts("CuSO4·5H2O");
473 let ascii = counts("CuSO4.5H2O");
474 assert_eq!(dot, ascii);
475 assert_eq!(dot, vec![(1, 10.0), (8, 9.0), (16, 1.0), (29, 1.0)]);
476 assert_counts("H2O.2H2O", &[(1, 6.0), (8, 3.0)]);
478 }
479
480 #[test]
481 fn tolerates_surrounding_whitespace_only() {
482 assert_eq!(counts(" H2O "), counts("H2O"));
483 }
484
485 #[test]
486 fn rejects_unbalanced_parentheses() {
487 let err = parse_formula("(H2O").unwrap_err();
488 assert!(
489 matches!(err, FormulaError::ParseError { pos: 4, .. }),
490 "{err}"
491 );
492 let err = parse_formula("H2O)").unwrap_err();
493 assert!(matches!(err, FormulaError::ParseError { .. }), "{err}");
494 }
495
496 #[test]
497 fn rejects_unknown_symbol() {
498 match parse_formula("Xx2O").unwrap_err() {
499 FormulaError::UnknownElement(s) => assert_eq!(s, "Xx"),
500 other => panic!("{other:?}"),
501 }
502 match parse_formula("Q").unwrap_err() {
503 FormulaError::UnknownElement(s) => assert_eq!(s, "Q"),
504 other => panic!("{other:?}"),
505 }
506 }
507
508 #[test]
509 fn rejects_leading_digit_empty_and_stray_chars() {
510 let err = parse_formula("2H2O").unwrap_err();
511 assert!(
512 matches!(err, FormulaError::ParseError { pos: 0, .. }),
513 "{err}"
514 );
515 assert!(matches!(
516 parse_formula("").unwrap_err(),
517 FormulaError::ParseError { .. }
518 ));
519 let err = parse_formula("H2 O").unwrap_err();
520 assert!(
521 matches!(err, FormulaError::ParseError { pos: 2, .. }),
522 "{err}"
523 );
524 let err = parse_formula("H1O-1").unwrap_err();
525 assert!(matches!(err, FormulaError::ParseError { .. }));
526 }
527
528 #[test]
529 fn from_formula_water_has_natural_isotopes_and_two_thirds_hydrogen() {
530 let water = Material::from_formula("H2O", &Ame2020, &NaturalAbundances, Some(1.0)).unwrap();
531 assert!(water
533 .comp
534 .contains_key(&NuclideId::from_name("H1").unwrap()));
535 for o in ["O16", "O17", "O18"] {
536 assert!(
537 water.comp.contains_key(&NuclideId::from_name(o).unwrap()),
538 "missing {o}"
539 );
540 }
541 let af = water.atom_fractions(&Ame2020).unwrap();
543 let h: f64 = af
544 .iter()
545 .filter(|(id, _)| id.z() == 1)
546 .map(|(_, f)| f)
547 .sum();
548 assert!((h - 2.0 / 3.0).abs() < 1e-12, "{h}");
549 let o: f64 = af
550 .iter()
551 .filter(|(id, _)| id.z() == 8)
552 .map(|(_, f)| f)
553 .sum();
554 assert!((o - 1.0 / 3.0).abs() < 1e-12);
555 assert_eq!(water.density(), Some(1.0));
557 }
558
559 #[test]
560 fn from_formula_h2so4() {
561 let mat = Material::from_formula("H2SO4", &Ame2020, &NaturalAbundances, None).unwrap();
562 let af = mat.atom_fractions(&Ame2020).unwrap();
563 assert!(!af.is_empty(), "H2SO4 atom fractions should not be empty");
564 for nuc in ["H1", "H2", "O16", "O17", "O18", "S32", "S33", "S34", "S36"] {
565 assert!(
566 af.contains_key(&NuclideId::from_name(nuc).unwrap()),
567 "missing {nuc}"
568 );
569 }
570 }
571
572 #[test]
573 fn from_formula_rejects_bad_input_and_missing_data() {
574 match Material::from_formula("Xx", &Ame2020, &NaturalAbundances, None).unwrap_err() {
575 FormulaError::UnknownElement(s) => assert_eq!(s, "Xx"),
576 other => panic!("{other:?}"),
577 }
578 assert!(matches!(
579 Material::from_formula("U", &Ame2020, &NoAbundances, None).unwrap_err(),
580 FormulaError::NoAbundanceData(92)
581 ));
582 }
583
584 #[test]
585 fn expand_then_collapse_round_trips_an_elemental_material() {
586 let mut mat = Material::new();
588 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();
591
592 mat.expand_elements(&Ame2020, &NaturalAbundances).unwrap();
593 assert!(!mat.comp.contains_key(&NuclideId::from_nucid(80_000_000)));
595 assert!(mat.comp.contains_key(&NuclideId::from_name("H1").unwrap()));
596 assert!(mat.comp.contains_key(&NuclideId::from_name("H2").unwrap()));
597 assert!(mat.comp.contains_key(&NuclideId::from_name("O18").unwrap()));
598
599 let back = mat.collapse_elements();
600 assert_eq!(
601 back.comp.keys().copied().collect::<Vec<_>>(),
602 original.comp.keys().copied().collect::<Vec<_>>()
603 );
604 for (id, m0) in &original.comp {
605 let m1 = back.comp[id];
606 assert!((m0 - m1).abs() < 1e-9 * m0.abs(), "{id}: {m0} vs {m1}");
607 }
608 }
609
610 #[test]
611 fn expand_preserves_entry_masses_and_leaves_named_nuclides_alone() {
612 let mut mat = Material::new();
613 mat.add_nuclide(NuclideId::from_nucid(10_000_000), 18.0); mat.add_nuclide(NuclideId::from_name("Fe56").unwrap(), 5.0);
615
616 mat.expand_elements(&Ame2020, &NaturalAbundances).unwrap();
617 close(mat.mass(), 23.0);
618 close(
619 mat.remove_nuclide(NuclideId::from_name("Fe56").unwrap())
620 .unwrap(),
621 5.0,
622 );
623 let h: f64 = mat.comp.values().sum();
624 close(h, 18.0);
625 }
626
627 #[test]
628 fn expand_without_abundances_errors_with_z() {
629 let mut mat = Material::new();
630 mat.add_nuclide(NuclideId::from_nucid(920_000_000), 1.0);
631 match mat.expand_elements(&Ame2020, &NoAbundances).unwrap_err() {
632 FormulaError::NoAbundanceData(z) => assert_eq!(z, 92),
633 other => panic!("{other:?}"),
634 }
635 }
636
637 #[test]
638 fn collapse_folds_named_nuclides_into_placeholder_keys() {
639 let mut mat = Material::new();
640 mat.add_nuclide(NuclideId::from_name("U235").unwrap(), 3.0);
641 mat.add_nuclide(NuclideId::from_name("U238").unwrap(), 1.0);
642 mat.set_density(Some(19.1));
643 let collapsed = mat.collapse_elements();
644
645 let key = NuclideId::from_nucid(920_000_000);
646 assert_eq!(collapsed.comp.len(), 1);
647 close(collapsed.comp[&key], 4.0);
648 assert_eq!(
650 key.nucid(),
651 nucleide_nuclei::element_z("U").unwrap() * 10_000_000
652 );
653 assert_eq!(collapsed.density(), Some(19.1));
654 }
655
656 fn close(a: f64, b: f64) {
657 assert!((a - b).abs() < 1e-12, "{a} != {b}");
658 }
659}