oxihuman_export/
geometry_nodes_export.rs

1//! Geometry node graph export (Blender-style procedural geometry graph description).
2
3#[allow(dead_code)]
4#[derive(Clone)]
5pub enum GeoNodeType {
6    MeshPrimitive,
7    Transform,
8    JoinGeometry,
9    SeparateGeometry,
10    Attribute,
11    Math,
12    Compare,
13    Switch,
14    Output,
15}
16
17impl GeoNodeType {
18    fn type_name(&self) -> &'static str {
19        match self {
20            GeoNodeType::MeshPrimitive => "MeshPrimitive",
21            GeoNodeType::Transform => "Transform",
22            GeoNodeType::JoinGeometry => "JoinGeometry",
23            GeoNodeType::SeparateGeometry => "SeparateGeometry",
24            GeoNodeType::Attribute => "Attribute",
25            GeoNodeType::Math => "Math",
26            GeoNodeType::Compare => "Compare",
27            GeoNodeType::Switch => "Switch",
28            GeoNodeType::Output => "Output",
29        }
30    }
31
32    fn discriminant(&self) -> u8 {
33        match self {
34            GeoNodeType::MeshPrimitive => 0,
35            GeoNodeType::Transform => 1,
36            GeoNodeType::JoinGeometry => 2,
37            GeoNodeType::SeparateGeometry => 3,
38            GeoNodeType::Attribute => 4,
39            GeoNodeType::Math => 5,
40            GeoNodeType::Compare => 6,
41            GeoNodeType::Switch => 7,
42            GeoNodeType::Output => 8,
43        }
44    }
45}
46
47#[allow(dead_code)]
48#[derive(Clone)]
49pub struct GeoNodeSocket {
50    pub name: String,
51    pub socket_type: String,
52    pub default_value: String,
53}
54
55#[allow(dead_code)]
56pub struct GeoNode {
57    pub id: u32,
58    pub name: String,
59    pub node_type: GeoNodeType,
60    pub inputs: Vec<GeoNodeSocket>,
61    pub outputs: Vec<GeoNodeSocket>,
62    pub position: [f32; 2],
63}
64
65#[allow(dead_code)]
66pub struct GeoNodeLink {
67    pub from_node: u32,
68    pub from_socket: usize,
69    pub to_node: u32,
70    pub to_socket: usize,
71}
72
73#[allow(dead_code)]
74pub struct GeoNodeGraph {
75    pub name: String,
76    pub nodes: Vec<GeoNode>,
77    pub links: Vec<GeoNodeLink>,
78    pub next_id: u32,
79}
80
81#[allow(dead_code)]
82pub fn new_geo_graph(name: &str) -> GeoNodeGraph {
83    GeoNodeGraph {
84        name: name.to_string(),
85        nodes: Vec::new(),
86        links: Vec::new(),
87        next_id: 1,
88    }
89}
90
91#[allow(dead_code)]
92pub fn add_geo_node(graph: &mut GeoNodeGraph, name: &str, node_type: GeoNodeType) -> u32 {
93    let id = graph.next_id;
94    graph.next_id += 1;
95    let x = id as f32 * 200.0;
96    graph.nodes.push(GeoNode {
97        id,
98        name: name.to_string(),
99        node_type,
100        inputs: Vec::new(),
101        outputs: Vec::new(),
102        position: [x, 0.0],
103    });
104    id
105}
106
107#[allow(dead_code)]
108pub fn add_geo_link(
109    graph: &mut GeoNodeGraph,
110    from: u32,
111    from_sock: usize,
112    to: u32,
113    to_sock: usize,
114) {
115    graph.links.push(GeoNodeLink {
116        from_node: from,
117        from_socket: from_sock,
118        to_node: to,
119        to_socket: to_sock,
120    });
121}
122
123#[allow(dead_code)]
124pub fn get_geo_node(graph: &GeoNodeGraph, id: u32) -> Option<&GeoNode> {
125    graph.nodes.iter().find(|n| n.id == id)
126}
127
128#[allow(dead_code)]
129pub fn geo_node_count(graph: &GeoNodeGraph) -> usize {
130    graph.nodes.len()
131}
132
133#[allow(dead_code)]
134pub fn geo_link_count(graph: &GeoNodeGraph) -> usize {
135    graph.links.len()
136}
137
138#[allow(dead_code)]
139pub fn export_geo_graph_json(graph: &GeoNodeGraph) -> String {
140    let nodes_json: Vec<String> = graph
141        .nodes
142        .iter()
143        .map(|n| {
144            format!(
145                "{{\"id\":{},\"name\":\"{}\",\"type\":\"{}\"}}",
146                n.id,
147                n.name,
148                n.node_type.type_name()
149            )
150        })
151        .collect();
152    let links_json: Vec<String> = graph
153        .links
154        .iter()
155        .map(|l| {
156            format!(
157                "{{\"from\":{},\"from_socket\":{},\"to\":{},\"to_socket\":{}}}",
158                l.from_node, l.from_socket, l.to_node, l.to_socket
159            )
160        })
161        .collect();
162    format!(
163        "{{\"name\":\"{}\",\"nodes\":[{}],\"links\":[{}]}}",
164        graph.name,
165        nodes_json.join(","),
166        links_json.join(",")
167    )
168}
169
170#[allow(dead_code)]
171pub fn export_geo_graph_python(graph: &GeoNodeGraph) -> String {
172    let mut lines = Vec::new();
173    lines.push("import bpy".to_string());
174    lines.push(format!("# Geometry node graph: {}", graph.name));
175    lines.push("node_tree = bpy.context.object.modifiers['GeometryNodes'].node_group".to_string());
176    lines.push("nodes = node_tree.nodes".to_string());
177    lines.push("links = node_tree.links".to_string());
178    lines.push("nodes.clear()".to_string());
179    lines.push(String::new());
180
181    let mut node_var_names: std::collections::HashMap<u32, String> =
182        std::collections::HashMap::new();
183    for node in &graph.nodes {
184        let var_name = format!("node_{}", node.id);
185        node_var_names.insert(node.id, var_name.clone());
186        let btype = match node.node_type {
187            GeoNodeType::MeshPrimitive => "GeometryNodeMeshCube",
188            GeoNodeType::Transform => "GeometryNodeTransform",
189            GeoNodeType::JoinGeometry => "GeometryNodeJoinGeometry",
190            GeoNodeType::SeparateGeometry => "GeometryNodeSeparateGeometry",
191            GeoNodeType::Attribute => "GeometryNodeInputNamedAttribute",
192            GeoNodeType::Math => "ShaderNodeMath",
193            GeoNodeType::Compare => "FunctionNodeCompare",
194            GeoNodeType::Switch => "GeometryNodeSwitch",
195            GeoNodeType::Output => "NodeGroupOutput",
196        };
197        lines.push(format!("{var_name} = nodes.new(type='{btype}')"));
198        lines.push(format!(
199            "{var_name}.location = ({}, {})",
200            node.position[0], node.position[1]
201        ));
202        lines.push(format!("{var_name}.label = \"{}\"", node.name));
203        lines.push(String::new());
204    }
205
206    for link in &graph.links {
207        if let (Some(from_var), Some(to_var)) = (
208            node_var_names.get(&link.from_node),
209            node_var_names.get(&link.to_node),
210        ) {
211            lines.push(format!(
212                "links.new({from_var}.outputs[{}], {to_var}.inputs[{}])",
213                link.from_socket, link.to_socket
214            ));
215        }
216    }
217
218    lines.join("\n")
219}
220
221#[allow(dead_code)]
222pub fn find_output_node(graph: &GeoNodeGraph) -> Option<&GeoNode> {
223    graph
224        .nodes
225        .iter()
226        .find(|n| n.node_type.discriminant() == GeoNodeType::Output.discriminant())
227}
228
229#[allow(dead_code)]
230pub fn nodes_of_type<'a>(graph: &'a GeoNodeGraph, node_type: &GeoNodeType) -> Vec<&'a GeoNode> {
231    let disc = node_type.discriminant();
232    graph
233        .nodes
234        .iter()
235        .filter(|n| n.node_type.discriminant() == disc)
236        .collect()
237}
238
239#[allow(dead_code)]
240pub fn remove_geo_node(graph: &mut GeoNodeGraph, id: u32) -> bool {
241    let before = graph.nodes.len();
242    graph.nodes.retain(|n| n.id != id);
243    graph.links.retain(|l| l.from_node != id && l.to_node != id);
244    graph.nodes.len() < before
245}
246
247#[allow(dead_code)]
248pub fn validate_geo_graph(graph: &GeoNodeGraph) -> Vec<String> {
249    let mut issues = Vec::new();
250    let node_ids: Vec<u32> = graph.nodes.iter().map(|n| n.id).collect();
251
252    for link in &graph.links {
253        if !node_ids.contains(&link.from_node) {
254            issues.push(format!(
255                "Dangling link: from_node {} not found",
256                link.from_node
257            ));
258        }
259        if !node_ids.contains(&link.to_node) {
260            issues.push(format!("Dangling link: to_node {} not found", link.to_node));
261        }
262    }
263
264    if find_output_node(graph).is_none() {
265        issues.push("No Output node found in graph".to_string());
266    }
267
268    issues
269}
270
271#[allow(dead_code)]
272pub fn default_output_node(graph: &mut GeoNodeGraph) -> u32 {
273    add_geo_node(graph, "Group Output", GeoNodeType::Output)
274}
275
276#[cfg(test)]
277mod tests {
278    use super::*;
279
280    #[test]
281    fn test_new_geo_graph() {
282        let graph = new_geo_graph("MyGraph");
283        assert_eq!(graph.name, "MyGraph");
284        assert!(graph.nodes.is_empty());
285        assert!(graph.links.is_empty());
286        assert_eq!(graph.next_id, 1);
287    }
288
289    #[test]
290    fn test_add_geo_node() {
291        let mut graph = new_geo_graph("G");
292        let id = add_geo_node(&mut graph, "Transform", GeoNodeType::Transform);
293        assert_eq!(id, 1);
294        assert_eq!(graph.nodes.len(), 1);
295        assert_eq!(graph.nodes[0].name, "Transform");
296    }
297
298    #[test]
299    fn test_add_geo_link() {
300        let mut graph = new_geo_graph("G");
301        let a = add_geo_node(&mut graph, "A", GeoNodeType::MeshPrimitive);
302        let b = add_geo_node(&mut graph, "B", GeoNodeType::Output);
303        add_geo_link(&mut graph, a, 0, b, 0);
304        assert_eq!(graph.links.len(), 1);
305        assert_eq!(graph.links[0].from_node, a);
306        assert_eq!(graph.links[0].to_node, b);
307    }
308
309    #[test]
310    fn test_get_geo_node_found() {
311        let mut graph = new_geo_graph("G");
312        let id = add_geo_node(&mut graph, "Math", GeoNodeType::Math);
313        let node = get_geo_node(&graph, id);
314        assert!(node.is_some());
315        assert_eq!(node.expect("should succeed").name, "Math");
316    }
317
318    #[test]
319    fn test_get_geo_node_not_found() {
320        let graph = new_geo_graph("G");
321        assert!(get_geo_node(&graph, 99).is_none());
322    }
323
324    #[test]
325    fn test_geo_node_count() {
326        let mut graph = new_geo_graph("G");
327        assert_eq!(geo_node_count(&graph), 0);
328        add_geo_node(&mut graph, "A", GeoNodeType::Math);
329        add_geo_node(&mut graph, "B", GeoNodeType::Math);
330        assert_eq!(geo_node_count(&graph), 2);
331    }
332
333    #[test]
334    fn test_geo_link_count() {
335        let mut graph = new_geo_graph("G");
336        let a = add_geo_node(&mut graph, "A", GeoNodeType::MeshPrimitive);
337        let b = add_geo_node(&mut graph, "B", GeoNodeType::Output);
338        assert_eq!(geo_link_count(&graph), 0);
339        add_geo_link(&mut graph, a, 0, b, 0);
340        assert_eq!(geo_link_count(&graph), 1);
341    }
342
343    #[test]
344    fn test_export_geo_graph_json_non_empty() {
345        let mut graph = new_geo_graph("TestGraph");
346        add_geo_node(&mut graph, "Node1", GeoNodeType::MeshPrimitive);
347        let json = export_geo_graph_json(&graph);
348        assert!(!json.is_empty());
349        assert!(json.contains("TestGraph"));
350        assert!(json.contains("MeshPrimitive"));
351    }
352
353    #[test]
354    fn test_export_geo_graph_python_non_empty() {
355        let mut graph = new_geo_graph("PyGraph");
356        add_geo_node(&mut graph, "Cube", GeoNodeType::MeshPrimitive);
357        let py = export_geo_graph_python(&graph);
358        assert!(!py.is_empty());
359        assert!(py.contains("import bpy"));
360    }
361
362    #[test]
363    fn test_find_output_node_none() {
364        let mut graph = new_geo_graph("G");
365        add_geo_node(&mut graph, "Math", GeoNodeType::Math);
366        assert!(find_output_node(&graph).is_none());
367    }
368
369    #[test]
370    fn test_find_output_node_found() {
371        let mut graph = new_geo_graph("G");
372        add_geo_node(&mut graph, "Output", GeoNodeType::Output);
373        let result = find_output_node(&graph);
374        assert!(result.is_some());
375    }
376
377    #[test]
378    fn test_remove_geo_node() {
379        let mut graph = new_geo_graph("G");
380        let a = add_geo_node(&mut graph, "A", GeoNodeType::Math);
381        let b = add_geo_node(&mut graph, "B", GeoNodeType::Output);
382        add_geo_link(&mut graph, a, 0, b, 0);
383        let removed = remove_geo_node(&mut graph, a);
384        assert!(removed);
385        assert_eq!(graph.nodes.len(), 1);
386        assert!(graph.links.is_empty());
387    }
388
389    #[test]
390    fn test_remove_geo_node_not_found() {
391        let mut graph = new_geo_graph("G");
392        add_geo_node(&mut graph, "A", GeoNodeType::Math);
393        let removed = remove_geo_node(&mut graph, 999);
394        assert!(!removed);
395    }
396
397    #[test]
398    fn test_validate_geo_graph_no_output_warning() {
399        let mut graph = new_geo_graph("G");
400        add_geo_node(&mut graph, "Math", GeoNodeType::Math);
401        let issues = validate_geo_graph(&graph);
402        assert!(!issues.is_empty());
403        assert!(issues.iter().any(|i| i.contains("Output")));
404    }
405
406    #[test]
407    fn test_validate_geo_graph_passes_with_output() {
408        let mut graph = new_geo_graph("G");
409        default_output_node(&mut graph);
410        let issues = validate_geo_graph(&graph);
411        assert!(issues.is_empty(), "Issues: {issues:?}");
412    }
413
414    #[test]
415    fn test_nodes_of_type() {
416        let mut graph = new_geo_graph("G");
417        add_geo_node(&mut graph, "M1", GeoNodeType::Math);
418        add_geo_node(&mut graph, "M2", GeoNodeType::Math);
419        add_geo_node(&mut graph, "Out", GeoNodeType::Output);
420        let math_nodes = nodes_of_type(&graph, &GeoNodeType::Math);
421        assert_eq!(math_nodes.len(), 2);
422    }
423
424    #[test]
425    fn test_default_output_node() {
426        let mut graph = new_geo_graph("G");
427        let id = default_output_node(&mut graph);
428        assert!(get_geo_node(&graph, id).is_some());
429        assert!(find_output_node(&graph).is_some());
430    }
431}
oxihuman_export/geometry_nodes_export.rs

oxihuman_export/
geometry_nodes_export.rs
