Skip to main content

xpict_core/
edge.rs

1//! EdgePlan / EdgeResult — host callback ABI (coord_gen / align).
2//!
3//! Core builds the plan (ids + forest); language edges process it and return
4//! flat [`MoleculeIn`] results. No molblocks on this wire.
5
6use serde::{Deserialize, Serialize};
7
8use crate::scene::MoleculeIn;
9
10#[cfg(feature = "codegen")]
11use schemars::JsonSchema;
12#[cfg(feature = "codegen")]
13use ts_rs::TS;
14
15/// Minimum mapped atoms before align is trusted (MCS or explicit map).
16pub const MIN_MCS_ATOMS: u32 = 3;
17
18/// Align this mol onto its parent template.
19#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
20#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
21#[cfg_attr(feature = "codegen", ts(export))]
22pub struct AlignOpts {
23    /// Pairs `(query_atom, template_atom)`. `None` → edge runs MCS.
24    #[serde(default, skip_serializing_if = "Option::is_none")]
25    #[cfg_attr(feature = "codegen", ts(optional))]
26    pub atom_map: Option<Vec<(u32, u32)>>,
27    /// Override [`MIN_MCS_ATOMS`] when set.
28    #[serde(default, skip_serializing_if = "Option::is_none")]
29    #[cfg_attr(feature = "codegen", ts(optional))]
30    pub min_atoms: Option<u32>,
31}
32
33/// One node in a coord_gen forest (root = free layout; children align to parent).
34#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
35#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
36#[cfg_attr(feature = "codegen", ts(export))]
37pub struct MolTemplate {
38    /// Rust-assigned unique id; round-trips to the document node.
39    pub id: String,
40    #[serde(default, skip_serializing_if = "Option::is_none")]
41    #[cfg_attr(feature = "codegen", ts(optional))]
42    pub smiles: Option<String>,
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    #[cfg_attr(feature = "codegen", ts(optional))]
45    pub cxsmiles: Option<String>,
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    #[cfg_attr(feature = "codegen", ts(optional))]
48    pub molfile: Option<String>,
49    /// Opts for aligning onto the parent; `None` on roots.
50    #[serde(default, skip_serializing_if = "Option::is_none")]
51    #[cfg_attr(feature = "codegen", ts(optional))]
52    pub align: Option<AlignOpts>,
53    /// Children that use this node as their align template.
54    #[serde(default, skip_serializing_if = "Vec::is_empty")]
55    pub template_for: Vec<MolTemplate>,
56}
57
58impl MolTemplate {
59    /// Exactly one of smiles / cxsmiles / molfile must be set.
60    pub fn validate_structure(&self) -> Result<(), String> {
61        let n = [&self.smiles, &self.cxsmiles, &self.molfile]
62            .iter()
63            .filter(|s| s.as_ref().is_some_and(|t| !t.trim().is_empty()))
64            .count();
65        if n == 1 {
66            Ok(())
67        } else {
68            Err(format!(
69                "MolTemplate {}: need exactly one of smiles/cxsmiles/molfile",
70                self.id
71            ))
72        }
73    }
74
75    pub fn source(&self) -> Result<&str, String> {
76        self.validate_structure()?;
77        if let Some(s) = self.smiles.as_deref().filter(|t| !t.trim().is_empty()) {
78            return Ok(s);
79        }
80        if let Some(s) = self.cxsmiles.as_deref().filter(|t| !t.trim().is_empty()) {
81            return Ok(s);
82        }
83        if let Some(s) = self.molfile.as_deref().filter(|t| !t.trim().is_empty()) {
84            return Ok(s);
85        }
86        Err(format!("MolTemplate {}: empty structure", self.id))
87    }
88}
89
90/// ``type: "coord_gen"`` edge task.
91#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
92#[serde(tag = "type", rename_all = "snake_case")]
93#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
94#[cfg_attr(feature = "codegen", ts(export))]
95pub enum EdgeTask {
96    CoordGen {
97        #[serde(default)]
98        roots: Vec<MolTemplate>,
99    },
100}
101
102/// Host callback request.
103#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
104#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
105#[cfg_attr(feature = "codegen", ts(export))]
106pub struct EdgePlan {
107    pub version: u32,
108    #[serde(default)]
109    pub tasks: Vec<EdgeTask>,
110}
111
112impl EdgePlan {
113    pub fn new_v1(tasks: Vec<EdgeTask>) -> Self {
114        Self {
115            version: 1,
116            tasks,
117        }
118    }
119
120    pub fn validate(&self) -> Result<(), String> {
121        if self.version != 1 {
122            return Err(format!("unsupported EdgePlan version {}", self.version));
123        }
124        let mut seen = std::collections::HashSet::new();
125        for task in &self.tasks {
126            match task {
127                EdgeTask::CoordGen { roots } => {
128                    for root in roots {
129                        validate_tree(root, /*is_root*/ true, &mut seen)?;
130                    }
131                }
132            }
133        }
134        Ok(())
135    }
136}
137
138fn validate_tree(
139    node: &MolTemplate,
140    is_root: bool,
141    seen: &mut std::collections::HashSet<String>,
142) -> Result<(), String> {
143    if !seen.insert(node.id.clone()) {
144        return Err(format!("duplicate MolTemplate id {}", node.id));
145    }
146    node.validate_structure()?;
147    if is_root && node.align.is_some() {
148        return Err(format!(
149            "MolTemplate {}: roots must have align=null",
150            node.id
151        ));
152    }
153    if !is_root && node.align.is_none() {
154        // Children may omit align object (= MCS defaults); that's OK.
155    }
156    for child in &node.template_for {
157        validate_tree(child, false, seen)?;
158    }
159    Ok(())
160}
161
162/// How coords were produced for one molecule.
163#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
164#[serde(rename_all = "snake_case")]
165#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
166#[cfg_attr(feature = "codegen", ts(export))]
167pub enum CoordMethod {
168    Free,
169    AtomMap,
170    Mcs,
171    /// Align was requested but failed — host fell back to unaligned coord gen.
172    None,
173}
174
175/// One flat molecule entry in a coord_gen result.
176#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
177#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
178#[cfg_attr(feature = "codegen", ts(export))]
179pub struct CoordGenMoleculeResult {
180    pub id: String,
181    /// True when usable coords were produced (aligned or free-layout fallback).
182    pub ok: bool,
183    pub method: CoordMethod,
184    #[serde(default, skip_serializing_if = "Option::is_none")]
185    #[cfg_attr(feature = "codegen", ts(optional))]
186    pub used_map: Option<Vec<(u32, u32)>>,
187    /// Present whenever ``ok`` is true.
188    #[serde(default, skip_serializing_if = "Option::is_none")]
189    #[cfg_attr(feature = "codegen", ts(optional))]
190    pub molecule: Option<MoleculeIn>,
191    #[serde(default, skip_serializing_if = "Option::is_none")]
192    #[cfg_attr(feature = "codegen", ts(optional))]
193    pub error: Option<String>,
194}
195
196#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
197#[serde(tag = "type", rename_all = "snake_case")]
198#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
199#[cfg_attr(feature = "codegen", ts(export))]
200pub enum EdgeTaskResult {
201    CoordGen {
202        ok: bool,
203        molecules: Vec<CoordGenMoleculeResult>,
204    },
205}
206
207#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
208#[cfg_attr(feature = "codegen", derive(JsonSchema, TS))]
209#[cfg_attr(feature = "codegen", ts(export))]
210pub struct EdgeResult {
211    pub version: u32,
212    #[serde(default)]
213    pub results: Vec<EdgeTaskResult>,
214}
215
216impl EdgeResult {
217    pub fn new_v1(results: Vec<EdgeTaskResult>) -> Self {
218        Self {
219            version: 1,
220            results,
221        }
222    }
223
224    /// Collect all molecule rows across tasks (order preserved).
225    pub fn molecules(&self) -> Vec<&CoordGenMoleculeResult> {
226        let mut out = Vec::new();
227        for r in &self.results {
228            match r {
229                EdgeTaskResult::CoordGen { molecules, .. } => out.extend(molecules.iter()),
230            }
231        }
232        out
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239
240    fn sample_plan() -> EdgePlan {
241        EdgePlan::new_v1(vec![EdgeTask::CoordGen {
242            roots: vec![MolTemplate {
243                id: "m_0".into(),
244                smiles: Some("c1ccccc1".into()),
245                cxsmiles: None,
246                molfile: None,
247                align: None,
248                template_for: vec![MolTemplate {
249                    id: "m_1".into(),
250                    smiles: Some("Cc1ccccc1".into()),
251                    cxsmiles: None,
252                    molfile: None,
253                    align: Some(AlignOpts {
254                        atom_map: Some(vec![
255                            (1, 0),
256                            (2, 1),
257                            (3, 2),
258                            (4, 3),
259                            (5, 4),
260                            (6, 5),
261                        ]),
262                        min_atoms: None,
263                    }),
264                    template_for: vec![],
265                }],
266            }],
267        }])
268    }
269
270    #[test]
271    fn plan_json_roundtrip() {
272        let plan = sample_plan();
273        plan.validate().unwrap();
274        let s = serde_json::to_string(&plan).unwrap();
275        let back: EdgePlan = serde_json::from_str(&s).unwrap();
276        assert_eq!(back, plan);
277    }
278
279    #[test]
280    fn reject_duplicate_ids() {
281        let mut plan = sample_plan();
282        match &mut plan.tasks[0] {
283            EdgeTask::CoordGen { roots } => {
284                roots[0].template_for[0].id = "m_0".into();
285            }
286        }
287        assert!(plan.validate().unwrap_err().contains("duplicate"));
288    }
289
290    #[test]
291    fn reject_root_with_align() {
292        let plan = EdgePlan::new_v1(vec![EdgeTask::CoordGen {
293            roots: vec![MolTemplate {
294                id: "m_0".into(),
295                smiles: Some("CCO".into()),
296                cxsmiles: None,
297                molfile: None,
298                align: Some(AlignOpts::default()),
299                template_for: vec![],
300            }],
301        }]);
302        assert!(plan.validate().unwrap_err().contains("align=null"));
303    }
304
305    #[test]
306    fn fallback_result_lists_all_ids() {
307        let result = EdgeResult::new_v1(vec![EdgeTaskResult::CoordGen {
308            ok: true,
309            molecules: vec![
310                CoordGenMoleculeResult {
311                    id: "m_0".into(),
312                    ok: true,
313                    method: CoordMethod::Free,
314                    used_map: None,
315                    molecule: None,
316                    error: None,
317                },
318                CoordGenMoleculeResult {
319                    id: "m_1".into(),
320                    ok: true,
321                    method: CoordMethod::None,
322                    used_map: None,
323                    molecule: None,
324                    error: Some("align failed; fell back".into()),
325                },
326            ],
327        }]);
328        let ids: Vec<_> = result.molecules().iter().map(|m| m.id.as_str()).collect();
329        assert_eq!(ids, ["m_0", "m_1"]);
330        assert!(matches!(
331            &result.results[0],
332            EdgeTaskResult::CoordGen { ok: true, .. }
333        ));
334    }
335}