use std::io::Read;
use std::sync::OnceLock;
use serde::{Deserialize, Serialize};
use super::cif::structure_from_cif;
use super::db::{
formula_matches, hit_from_structure, StructureHit, StructureQuery, StructureSource,
};
use super::model::Structure;
use super::StructureError;
static BLOB: &[u8] = include_bytes!("builtin_library.json.gz");
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct BuiltinEntry {
pub key: String,
pub name: String,
pub formula: String,
pub sg_number: u16,
#[serde(default)]
pub sg: Option<String>,
pub category: String,
pub source: String,
pub id: String,
#[serde(default)]
pub url: Option<String>,
#[serde(default)]
pub year: Option<String>,
#[serde(default)]
pub journal: Option<String>,
#[serde(default)]
pub authors: Option<String>,
#[serde(default)]
pub doi: Option<String>,
pub license: String,
pub citation: String,
#[serde(default)]
pub note: String,
pub cif: String,
}
#[derive(Debug)]
pub struct BuiltinLibrary {
entries: Vec<BuiltinEntry>,
hits: Vec<StructureHit>,
}
static LIBRARY: OnceLock<Result<BuiltinLibrary, String>> = OnceLock::new();
impl BuiltinLibrary {
pub fn get() -> Result<&'static BuiltinLibrary, StructureError> {
LIBRARY
.get_or_init(|| Self::decode(BLOB).map_err(|e| e.to_string()))
.as_ref()
.map_err(|e| StructureError::Database {
reason: format!("built-in library: {e}"),
})
}
fn decode(blob: &[u8]) -> Result<Self, StructureError> {
let mut text = String::new();
flate2::read::GzDecoder::new(blob)
.read_to_string(&mut text)
.map_err(|e| StructureError::Database {
reason: format!("built-in library is not valid gzip: {e}"),
})?;
let entries: Vec<BuiltinEntry> = serde_json::from_str(&text)?;
let hits = entries
.iter()
.map(|e| {
let mut hit = StructureHit {
id: e.key.clone(),
source: "builtin".into(),
formula: e.formula.replace(' ', ""),
name: Some(e.name.clone()),
space_group: e.sg.clone(),
elements: elements_of(&e.formula),
extra: Default::default(),
};
hit.extra.insert("category".into(), e.category.clone());
hit.extra
.insert("origin".into(), format!("{} {}", e.source, e.id));
if let Some(u) = &e.url {
hit.extra.insert("url".into(), u.clone());
}
hit.extra.insert("license".into(), e.license.clone());
hit
})
.collect();
Ok(Self { entries, hits })
}
pub fn entries(&self) -> &[BuiltinEntry] {
&self.entries
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn entry(&self, key: &str) -> Option<&BuiltinEntry> {
self.entries.iter().find(|e| e.key == key)
}
pub fn structure(&self, key: &str) -> Result<Structure, StructureError> {
let entry = self.entry(key).ok_or_else(|| StructureError::Database {
reason: format!("no built-in structure '{key}'"),
})?;
let mut s = structure_from_cif(&entry.cif)?;
s.title = entry.name.clone();
s.source = format!("built-in ({} {})", entry.source, entry.id);
if s.mineral.is_none() {
s.mineral = Some(entry.name.clone());
}
Ok(s)
}
}
fn elements_of(formula: &str) -> Vec<String> {
formula
.split_whitespace()
.map(|tok| {
tok.chars()
.take_while(|c| c.is_ascii_alphabetic())
.collect::<String>()
})
.filter(|s| !s.is_empty())
.collect()
}
impl StructureSource for BuiltinLibrary {
fn name(&self) -> &str {
"built-in"
}
fn search(&self, query: &StructureQuery) -> Result<Vec<StructureHit>, StructureError> {
let text = query.text.as_deref().map(str::trim).unwrap_or("");
let elements_only = StructureQuery {
text: None,
..query.clone()
};
let mut hits: Vec<StructureHit> = self
.hits
.iter()
.filter(|h| elements_only.matches(h) && (query.matches(h) || name_matches(h, text)))
.cloned()
.collect();
let lower = text.to_ascii_lowercase();
let exact = |h: &StructureHit| {
!lower.is_empty() && (h.id == lower || formula_matches(&h.formula, text))
};
hits.sort_by(|a, b| exact(b).cmp(&exact(a)).then_with(|| a.name.cmp(&b.name)));
if query.limit > 0 {
hits.truncate(query.limit);
}
Ok(hits)
}
fn fetch(&self, hit: &StructureHit) -> Result<Structure, StructureError> {
let mut s = self.structure(&hit.id)?;
let h = hit_from_structure(&s, "builtin", &hit.id);
if s.formula_sum.is_none() {
s.formula_sum = Some(h.formula);
}
Ok(s)
}
}
fn name_matches(hit: &StructureHit, text: &str) -> bool {
if text.is_empty() {
return true;
}
let t = text.to_ascii_lowercase();
hit.name
.as_deref()
.is_some_and(|n| n.to_ascii_lowercase().contains(&t))
|| hit.id.contains(&t)
|| hit
.extra
.get("category")
.is_some_and(|c| c.eq_ignore_ascii_case(&t))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::xafs::structure::cluster::{
absorber_sites, build_cluster, AbsorberSelection, ClusterOptions,
};
#[test]
fn library_decodes_and_every_entry_expands_and_clusters() {
let lib = BuiltinLibrary::get().expect("bundled library");
assert!(lib.len() >= 40, "only {} entries", lib.len());
for e in lib.entries() {
let s = lib
.structure(&e.key)
.unwrap_or_else(|err| panic!("{}: {err}", e.key));
assert!(!s.sites.is_empty(), "{} has no sites", e.key);
assert_eq!(
s.space_group.number,
Some(e.sg_number),
"{}: space group {:?} vs catalog {}",
e.key,
s.space_group.number,
e.sg_number
);
let first = s.elements()[0].clone();
let sites = absorber_sites(&s, &first);
assert!(!sites.is_empty(), "{}: no {first} site", e.key);
let cluster = build_cluster(
&s,
&AbsorberSelection::SiteIndex(sites[0]),
&ClusterOptions {
radius: 8.0,
..Default::default()
},
)
.unwrap_or_else(|err| panic!("{}: cluster {err}", e.key));
assert!(
cluster.atoms.len() > 10,
"{}: only {} atoms",
e.key,
cluster.atoms.len()
);
assert!(cluster.shells(0.01).len() >= 2, "{}: too few shells", e.key);
}
}
#[test]
fn search_by_name_formula_and_element() {
let lib = BuiltinLibrary::get().unwrap();
let hits = lib.search(&StructureQuery::text("ruthenium")).unwrap();
assert!(hits.iter().any(|h| h.id == "ru_hcp"), "{hits:?}");
let hits = lib.search(&StructureQuery::text("RuO2")).unwrap();
assert_eq!(hits.first().map(|h| h.id.as_str()), Some("ruo2"));
let hits = lib
.search(&StructureQuery::text("").with_elements(["Fe", "S"]))
.unwrap();
assert!(hits.iter().any(|h| h.id == "fes2_pyrite"));
let all = lib.search(&StructureQuery::text("")).unwrap();
assert_eq!(all.len(), lib.len());
let s = lib.fetch(&hits[0]).unwrap();
assert!(
s.elements().iter().any(|e| e == "Fe"),
"{} -> {:?} {:?}",
hits[0].id,
s.title,
s.elements()
);
}
}