nir_rs/
graph.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! In-memory NIR graph model.
4//!
5//! A graph is a named set of computational nodes plus a list of directed
6//! identity edges, matching neuromorphs/NIR (`nodes`, `edges`, `metadata`,
7//! optional `version`). Cycles are allowed; structure validation only checks
8//! edge endpoints and duplicate directed edges.
9
10use crate::error::{NirError, Result};
11use crate::nodes::NirNode;
12use crate::types::MetadataValue;
13use indexmap::IndexMap;
14use std::collections::{HashMap, HashSet};
15
16/// Directed NIR computation graph.
17///
18/// Node insertion order is preserved via [`IndexMap`] (stable iteration for
19/// serialization and debugging). Edges are an ordered list of `(src, dst)`
20/// name pairs.
21#[derive(Debug, Clone, PartialEq, Default)]
22#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
23pub struct NirGraph {
24    /// Named computational nodes (insertion-ordered).
25    pub nodes: IndexMap<String, NirNode>,
26    /// Directed edges as `(source_name, destination_name)`.
27    pub edges: Vec<(String, String)>,
28    /// Free-form graph metadata.
29    pub metadata: HashMap<String, MetadataValue>,
30    /// Optional NIR version string (set when loading from HDF5 in v0.3).
31    pub version: Option<String>,
32}
33
34impl NirGraph {
35    /// Create an empty graph.
36    #[must_use]
37    pub fn new() -> Self {
38        Self::default()
39    }
40
41    /// Insert a node under `name`.
42    ///
43    /// # Errors
44    ///
45    /// Returns [`NirError::DuplicateNode`] if `name` is already present.
46    pub fn insert_node(&mut self, name: impl Into<String>, node: NirNode) -> Result<()> {
47        let name = name.into();
48        if self.nodes.contains_key(&name) {
49            return Err(NirError::DuplicateNode(name));
50        }
51        self.nodes.insert(name, node);
52        Ok(())
53    }
54
55    /// Append a directed edge `(from → to)` without validating endpoints.
56    ///
57    /// Call [`validate_structure`](Self::validate_structure) to check that
58    /// endpoints exist and that the edge is not duplicated.
59    pub fn add_edge(&mut self, from: impl Into<String>, to: impl Into<String>) {
60        self.edges.push((from.into(), to.into()));
61    }
62
63    /// Borrow the node named `name`, if present.
64    #[must_use]
65    pub fn get(&self, name: &str) -> Option<&NirNode> {
66        self.nodes.get(name)
67    }
68
69    /// Mutably borrow the node named `name`, if present.
70    pub fn get_mut(&mut self, name: &str) -> Option<&mut NirNode> {
71        self.nodes.get_mut(name)
72    }
73
74    /// Number of nodes.
75    #[must_use]
76    pub fn len(&self) -> usize {
77        self.nodes.len()
78    }
79
80    /// Whether the graph has no nodes.
81    #[must_use]
82    pub fn is_empty(&self) -> bool {
83        self.nodes.is_empty()
84    }
85
86    /// Validate structural integrity.
87    ///
88    /// Checks:
89    /// - every edge endpoint names an existing node
90    /// - no duplicate directed edges `(src, dst)`
91    /// - nested [`NirNode::Graph`] subgraphs also validate
92    ///
93    /// Cycles are **allowed**. Type/shape inference is out of scope for v0.2.
94    ///
95    /// # Errors
96    ///
97    /// - [`NirError::MissingNode`] if an endpoint is unknown
98    /// - [`NirError::DuplicateEdge`] if the same directed edge appears twice
99    /// - [`NirError::InvalidGraph`] if a nested subgraph fails validation
100    pub fn validate_structure(&self) -> Result<()> {
101        let node_keys: HashSet<&str> = self.nodes.keys().map(String::as_str).collect();
102
103        for (src, dst) in &self.edges {
104            if !node_keys.contains(src.as_str()) {
105                return Err(NirError::MissingNode(src.clone()));
106            }
107            if !node_keys.contains(dst.as_str()) {
108                return Err(NirError::MissingNode(dst.clone()));
109            }
110        }
111
112        let mut seen_edges: HashSet<(&str, &str)> = HashSet::new();
113        for (src, dst) in &self.edges {
114            let key = (src.as_str(), dst.as_str());
115            if !seen_edges.insert(key) {
116                return Err(NirError::DuplicateEdge(src.clone(), dst.clone()));
117            }
118        }
119
120        // Nested graphs: re-prefix structure errors so callers see which subgraph failed.
121        // Today this method only returns MissingNode / DuplicateEdge / InvalidGraph;
122        // the `other` arm preserves any future validation variants as InvalidGraph.
123        for (name, node) in &self.nodes {
124            if let NirNode::Graph(sub) = node {
125                sub.validate_structure().map_err(|e| match e {
126                    NirError::MissingNode(n) => {
127                        NirError::InvalidGraph(format!("in subgraph {name:?}: missing node: {n}"))
128                    }
129                    NirError::DuplicateEdge(a, b) => NirError::InvalidGraph(format!(
130                        "in subgraph {name:?}: duplicate edge: ({a}, {b})"
131                    )),
132                    NirError::DuplicateNode(n) => {
133                        NirError::InvalidGraph(format!("in subgraph {name:?}: duplicate node: {n}"))
134                    }
135                    NirError::InvalidGraph(msg) => {
136                        NirError::InvalidGraph(format!("in subgraph {name:?}: {msg}"))
137                    }
138                    other => NirError::InvalidGraph(format!("in subgraph {name:?}: {other}")),
139                })?;
140            }
141        }
142
143        Ok(())
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150    use crate::nodes::{Affine, Input, Lif, Output};
151    use crate::types::Tensor;
152
153    fn input(shape: Vec<usize>) -> NirNode {
154        NirNode::Input(Input {
155            shape,
156            metadata: Default::default(),
157        })
158    }
159
160    fn output(shape: Vec<usize>) -> NirNode {
161        NirNode::Output(Output {
162            shape,
163            metadata: Default::default(),
164        })
165    }
166
167    #[test]
168    fn empty_graph_default() {
169        let g = NirGraph::new();
170        assert!(g.is_empty());
171        assert_eq!(g.len(), 0);
172        assert!(g.validate_structure().is_ok());
173    }
174
175    #[test]
176    fn insert_node_rejects_duplicate() {
177        let mut g = NirGraph::new();
178        g.insert_node("a", input(vec![4])).unwrap();
179        let err = g.insert_node("a", input(vec![2])).unwrap_err();
180        assert_eq!(err, NirError::DuplicateNode("a".into()));
181    }
182
183    #[test]
184    fn get_returns_inserted_node() {
185        let mut g = NirGraph::new();
186        g.insert_node("in", input(vec![3])).unwrap();
187        assert!(matches!(g.get("in"), Some(NirNode::Input(_))));
188        assert!(g.get("missing").is_none());
189    }
190
191    #[test]
192    fn validate_missing_source_endpoint() {
193        let mut g = NirGraph::new();
194        g.insert_node("b", output(vec![1])).unwrap();
195        g.add_edge("ghost", "b");
196        let err = g.validate_structure().unwrap_err();
197        assert_eq!(err, NirError::MissingNode("ghost".into()));
198    }
199
200    #[test]
201    fn validate_missing_dest_endpoint() {
202        let mut g = NirGraph::new();
203        g.insert_node("a", input(vec![1])).unwrap();
204        g.add_edge("a", "ghost");
205        let err = g.validate_structure().unwrap_err();
206        assert_eq!(err, NirError::MissingNode("ghost".into()));
207    }
208
209    #[test]
210    fn validate_duplicate_edge() {
211        let mut g = NirGraph::new();
212        g.insert_node("a", input(vec![1])).unwrap();
213        g.insert_node("b", output(vec![1])).unwrap();
214        g.add_edge("a", "b");
215        g.add_edge("a", "b");
216        let err = g.validate_structure().unwrap_err();
217        assert_eq!(err, NirError::DuplicateEdge("a".into(), "b".into()));
218    }
219
220    #[test]
221    fn cycles_are_allowed() {
222        let mut g = NirGraph::new();
223        g.insert_node("a", input(vec![1])).unwrap();
224        g.insert_node("b", output(vec![1])).unwrap();
225        g.add_edge("a", "b");
226        g.add_edge("b", "a");
227        assert!(g.validate_structure().is_ok());
228    }
229
230    #[test]
231    fn integration_input_affine_lif_output() {
232        let weight = Tensor::from_f32(vec![2, 4], vec![0.1; 8]).unwrap();
233        let bias = Tensor::from_f32(vec![2], vec![0.0, 0.0]).unwrap();
234        let tau = Tensor::from_f64(vec![2], vec![10.0, 10.0]).unwrap();
235        let r = Tensor::from_f64(vec![2], vec![1.0, 1.0]).unwrap();
236        let v_leak = Tensor::from_f64(vec![2], vec![0.0, 0.0]).unwrap();
237        let v_th = Tensor::from_f64(vec![2], vec![1.0, 1.0]).unwrap();
238
239        let mut g = NirGraph::new();
240        g.version = Some("1.0.0".into());
241        g.metadata.insert(
242            "origin".into(),
243            MetadataValue::String("integration-test".into()),
244        );
245
246        g.insert_node("input", input(vec![4])).unwrap();
247        g.insert_node(
248            "fc",
249            NirNode::Affine(Affine {
250                weight,
251                bias,
252                metadata: Default::default(),
253            }),
254        )
255        .unwrap();
256        g.insert_node(
257            "lif",
258            NirNode::Lif(Lif {
259                tau,
260                r,
261                v_leak,
262                v_threshold: v_th,
263                v_reset: None,
264                metadata: Default::default(),
265            }),
266        )
267        .unwrap();
268        g.insert_node("output", output(vec![2])).unwrap();
269
270        g.add_edge("input", "fc");
271        g.add_edge("fc", "lif");
272        g.add_edge("lif", "output");
273
274        assert!(g.validate_structure().is_ok());
275        assert_eq!(g.len(), 4);
276        assert_eq!(g.edges.len(), 3);
277        assert_eq!(g.get("lif").unwrap().type_name(), "LIF");
278        assert_eq!(g.get("fc").unwrap().type_name(), "Affine");
279    }
280
281    #[test]
282    fn nested_graph_validates() {
283        let mut inner = NirGraph::new();
284        inner.insert_node("i", input(vec![1])).unwrap();
285        inner.insert_node("o", output(vec![1])).unwrap();
286        inner.add_edge("i", "o");
287
288        let mut outer = NirGraph::new();
289        outer
290            .insert_node("sub", NirNode::Graph(Box::new(inner)))
291            .unwrap();
292        outer.insert_node("out", output(vec![1])).unwrap();
293        outer.add_edge("sub", "out");
294        assert!(outer.validate_structure().is_ok());
295    }
296
297    #[test]
298    fn nested_graph_reports_inner_failure() {
299        let mut inner = NirGraph::new();
300        inner.insert_node("i", input(vec![1])).unwrap();
301        // edge to missing node
302        inner.add_edge("i", "missing");
303
304        let mut outer = NirGraph::new();
305        outer
306            .insert_node("sub", NirNode::Graph(Box::new(inner)))
307            .unwrap();
308        let err = outer.validate_structure().unwrap_err();
309        match err {
310            NirError::InvalidGraph(msg) => {
311                assert!(msg.contains("sub"));
312                assert!(msg.contains("missing"));
313            }
314            other => panic!("expected InvalidGraph, got {other:?}"),
315        }
316    }
317
318    #[test]
319    fn insertion_order_preserved() {
320        let mut g = NirGraph::new();
321        g.insert_node("z", input(vec![1])).unwrap();
322        g.insert_node("a", output(vec![1])).unwrap();
323        let keys: Vec<&str> = g.nodes.keys().map(String::as_str).collect();
324        assert_eq!(keys, ["z", "a"]);
325    }
326}
nir_rs/graph.rs

nir_rs/
graph.rs
