Skip to main content

onnx_runtime_ir/
graph.rs

1//! The mutable graph model: node/value arenas, edge-consistent mutation,
2//! topological ordering, and validation (see `docs/ORT2.md` §3.3 and §3.5).
3
4use std::collections::{BinaryHeap, HashMap, HashSet};
5
6use crate::arena::Arena;
7use crate::dtype::DataType;
8use crate::error::GraphError;
9use crate::node::{Node, NodeId};
10use crate::shape::{Shape, SymbolConstraints, SymbolId};
11use crate::tensor::WeightRef;
12use crate::value::{Value, ValueId};
13
14/// A computation graph in SSA form.
15///
16/// Nodes and values live in [`Arena`]s keyed by [`NodeId`] / [`ValueId`]. The
17/// mutation API keeps producer/consumer edges consistent, so optimization
18/// passes can rewrite the graph and then [`Graph::validate`] it.
19#[derive(Clone, Debug, Default)]
20pub struct Graph {
21    pub nodes: Arena<NodeId, Node>,
22    pub values: Arena<ValueId, Value>,
23    /// Graph inputs, in order. These have no producer.
24    pub inputs: Vec<ValueId>,
25    /// Graph outputs, in order.
26    pub outputs: Vec<ValueId>,
27    /// Constant initializer weights, keyed by the value they populate.
28    pub initializers: HashMap<ValueId, WeightRef>,
29    /// Constraints on symbolic dimensions.
30    pub symbol_constraints: HashMap<SymbolId, SymbolConstraints>,
31    /// Imported opsets: domain → version.
32    pub opset_imports: HashMap<String, u64>,
33    /// Subgraph bodies for control-flow ops, keyed by `(node, attr_name)`.
34    pub subgraphs: HashMap<(NodeId, String), Graph>,
35
36    next_symbol: u32,
37    symbol_names: HashMap<String, SymbolId>,
38    unknown_value_types: HashSet<ValueId>,
39    unknown_value_shapes: HashSet<ValueId>,
40}
41
42impl Graph {
43    /// An empty graph.
44    pub fn new() -> Self {
45        Self::default()
46    }
47
48    // === Query API ===
49
50    /// Borrow a node. Panics if `id` is not live; use
51    /// [`Graph::try_node`] for a checked lookup.
52    pub fn node(&self, id: NodeId) -> &Node {
53        self.nodes.get(id).expect("node id not live in graph")
54    }
55
56    /// Mutably borrow a node. Panics if `id` is not live.
57    pub fn node_mut(&mut self, id: NodeId) -> &mut Node {
58        self.nodes.get_mut(id).expect("node id not live in graph")
59    }
60
61    /// Checked node lookup.
62    pub fn try_node(&self, id: NodeId) -> Option<&Node> {
63        self.nodes.get(id)
64    }
65
66    /// Borrow a value. Panics if `id` is not live; use
67    /// [`Graph::try_value`] for a checked lookup.
68    pub fn value(&self, id: ValueId) -> &Value {
69        self.values.get(id).expect("value id not live in graph")
70    }
71
72    /// Mutably borrow a value. Panics if `id` is not live.
73    pub fn value_mut(&mut self, id: ValueId) -> &mut Value {
74        self.values.get_mut(id).expect("value id not live in graph")
75    }
76
77    /// Checked value lookup.
78    pub fn try_value(&self, id: ValueId) -> Option<&Value> {
79        self.values.get(id)
80    }
81
82    /// Consuming input slots sorted by `(NodeId, input_index)`.
83    pub fn uses(&self, value: ValueId) -> Vec<(NodeId, u32)> {
84        self.value(value).consumers.uses()
85    }
86
87    /// Distinct consumer nodes sorted by ascending [`NodeId`].
88    pub fn consumers(&self, value: ValueId) -> Vec<NodeId> {
89        self.value(value).consumers.nodes()
90    }
91
92    /// Number of consuming input slots.
93    pub fn num_uses(&self, value: ValueId) -> usize {
94        self.value(value).consumers.len()
95    }
96
97    /// Whether at least one node input slot consumes `value`.
98    pub fn has_uses(&self, value: ValueId) -> bool {
99        !self.value(value).consumers.is_empty()
100    }
101
102    /// Number of live nodes.
103    pub fn num_nodes(&self) -> usize {
104        self.nodes.len()
105    }
106
107    /// Number of live values.
108    pub fn num_values(&self) -> usize {
109        self.values.len()
110    }
111
112    /// Whether a value's element type came from explicit source type information.
113    pub fn value_type_is_known(&self, id: ValueId) -> bool {
114        !self.unknown_value_types.contains(&id)
115    }
116
117    /// Whether a value's rank and dimensions came from explicit source shape information.
118    pub fn value_shape_is_known(&self, id: ValueId) -> bool {
119        !self.unknown_value_shapes.contains(&id)
120    }
121
122    /// Mark a value's placeholder element type as unknown.
123    pub fn mark_value_type_unknown(&mut self, id: ValueId) {
124        self.unknown_value_types.insert(id);
125    }
126
127    /// Mark a value's placeholder shape as unknown.
128    pub fn mark_value_shape_unknown(&mut self, id: ValueId) {
129        self.unknown_value_shapes.insert(id);
130    }
131
132    // === Symbolic dimensions ===
133
134    /// Allocate a fresh symbolic dimension with an optional name (no dedup).
135    pub fn create_symbol(&mut self, name: Option<String>) -> SymbolId {
136        let id = SymbolId(self.next_symbol);
137        self.next_symbol += 1;
138        self.symbol_constraints
139            .insert(id, SymbolConstraints::new(id, name.clone()));
140        if let Some(n) = name {
141            self.symbol_names.insert(n, id);
142        }
143        id
144    }
145
146    /// Intern a symbolic dimension by protobuf dim-param name: repeated names
147    /// resolve to the same [`SymbolId`] (graph-construction invariant §3.5.4).
148    pub fn intern_symbol(&mut self, name: &str) -> SymbolId {
149        if let Some(id) = self.symbol_names.get(name) {
150            return *id;
151        }
152        self.create_symbol(Some(name.to_string()))
153    }
154
155    // === Construction helpers ===
156
157    /// Create a new anonymous value with a contiguous default layout.
158    pub fn create_value(&mut self, dtype: DataType, shape: Shape) -> ValueId {
159        self.values
160            .insert_with(|vid| Value::new(vid, dtype, shape.clone()))
161    }
162
163    /// Create a new named value.
164    pub fn create_named_value(
165        &mut self,
166        name: impl Into<String>,
167        dtype: DataType,
168        shape: Shape,
169    ) -> ValueId {
170        let id = self.create_value(dtype, shape);
171        self.value_mut(id).name = Some(name.into());
172        id
173    }
174
175    /// Register `value` as a graph input.
176    pub fn add_input(&mut self, value: ValueId) {
177        self.value_mut(value).is_graph_input = true;
178        self.inputs.push(value);
179        debug_assert!(self.value(value).is_graph_input);
180    }
181
182    /// Register `value` as a graph output.
183    pub fn add_output(&mut self, value: ValueId) {
184        self.value_mut(value).is_graph_output = true;
185        self.outputs.push(value);
186        debug_assert!(self.value(value).is_graph_output);
187    }
188
189    /// Insert `value` into the ordered graph outputs.
190    pub fn insert_output(&mut self, index: usize, value: ValueId) {
191        self.value_mut(value).is_graph_output = true;
192        self.outputs.insert(index, value);
193        debug_assert!(self.value(value).is_graph_output);
194    }
195
196    /// Remove one ordered graph input.
197    pub fn remove_input(&mut self, index: usize) -> ValueId {
198        let value = self.inputs.remove(index);
199        if !self.inputs.contains(&value) {
200            self.value_mut(value).is_graph_input = false;
201        }
202        debug_assert_eq!(
203            self.value(value).is_graph_input,
204            self.inputs.contains(&value)
205        );
206        value
207    }
208
209    /// Remove one ordered graph output.
210    pub fn remove_output(&mut self, index: usize) -> ValueId {
211        let value = self.outputs.remove(index);
212        if !self.outputs.contains(&value) {
213            self.value_mut(value).is_graph_output = false;
214        }
215        debug_assert_eq!(
216            self.value(value).is_graph_output,
217            self.outputs.contains(&value)
218        );
219        value
220    }
221
222    /// Replace the complete ordered graph-input list.
223    pub fn set_inputs(&mut self, inputs: Vec<ValueId>) {
224        for value in self.inputs.drain(..) {
225            if let Some(metadata) = self.values.get_mut(value) {
226                metadata.is_graph_input = false;
227            }
228        }
229        for &value in &inputs {
230            self.value_mut(value).is_graph_input = true;
231        }
232        self.inputs = inputs;
233        debug_assert!(
234            self.inputs
235                .iter()
236                .all(|&value| self.value(value).is_graph_input)
237        );
238    }
239
240    /// Replace the complete ordered graph-output list.
241    pub fn set_outputs(&mut self, outputs: Vec<ValueId>) {
242        for value in self.outputs.drain(..) {
243            if let Some(metadata) = self.values.get_mut(value) {
244                metadata.is_graph_output = false;
245            }
246        }
247        for &value in &outputs {
248            self.value_mut(value).is_graph_output = true;
249        }
250        self.outputs = outputs;
251        debug_assert!(
252            self.outputs
253                .iter()
254                .all(|&value| self.value(value).is_graph_output)
255        );
256    }
257
258    /// Attach initializer weights to `value`.
259    pub fn set_initializer(&mut self, value: ValueId, weight: WeightRef) {
260        self.initializers.insert(value, weight);
261    }
262
263    // === Traversal ===
264
265    /// Direct predecessors: nodes that produce this node's inputs.
266    pub fn predecessors(&self, node: NodeId) -> Vec<NodeId> {
267        let mut out = Vec::new();
268        let mut seen = HashSet::new();
269        for v in self.node(node).input_values() {
270            if let Some(val) = self.values.get(v)
271                && let Some(prod) = val.producer
272                && seen.insert(prod)
273            {
274                out.push(prod);
275            }
276        }
277        out.sort_unstable_by_key(|node| node.0);
278        out
279    }
280
281    /// Direct successors: nodes that consume this node's outputs.
282    pub fn successors(&self, node: NodeId) -> Vec<NodeId> {
283        let mut out = Vec::new();
284        let mut seen = HashSet::new();
285        for &v in &self.node(node).outputs {
286            if let Some(val) = self.values.get(v) {
287                for c in val.consumers.nodes() {
288                    if seen.insert(c) {
289                        out.push(c);
290                    }
291                }
292            }
293        }
294        out.sort_unstable_by_key(|node| node.0);
295        out
296    }
297
298    /// All nodes that lie on a path between `inputs` and `outputs`.
299    ///
300    /// Walks backwards from `outputs` via producer edges, stopping at any value
301    /// in `inputs`. Used to extract subgraphs for EP capability claims (§3.4).
302    pub fn nodes_between(&self, inputs: &[ValueId], outputs: &[ValueId]) -> Vec<NodeId> {
303        let boundary: HashSet<ValueId> = inputs.iter().copied().collect();
304        let mut nodes = Vec::new();
305        let mut seen_nodes = HashSet::new();
306        let mut seen_values = HashSet::new();
307        let mut stack: Vec<ValueId> = outputs.to_vec();
308        while let Some(v) = stack.pop() {
309            if boundary.contains(&v) || !seen_values.insert(v) {
310                continue;
311            }
312            let Some(val) = self.values.get(v) else {
313                continue;
314            };
315            if let Some(prod) = val.producer {
316                if seen_nodes.insert(prod) {
317                    nodes.push(prod);
318                }
319                for iv in self.node(prod).input_values() {
320                    stack.push(iv);
321                }
322            }
323        }
324        nodes
325    }
326
327    /// Topological order of nodes via Kahn's algorithm.
328    ///
329    /// Ties are broken by ascending [`NodeId`] for deterministic output.
330    /// Returns [`GraphError::CycleDetected`] if the graph has a cycle.
331    pub fn topological_order(&self) -> Result<Vec<NodeId>, GraphError> {
332        const VACANT: usize = usize::MAX;
333        let mut in_degree = vec![VACANT; self.nodes.capacity()];
334        let mut adj = vec![Vec::<NodeId>::new(); self.nodes.capacity()];
335        for node in self.nodes.keys() {
336            in_degree[node.0 as usize] = 0;
337        }
338
339        for (nid, node) in self.nodes.iter() {
340            for v in node.input_values() {
341                if let Some(val) = self.values.get(v)
342                    && let Some(prod) = val.producer
343                    && self.nodes.contains(prod)
344                {
345                    adj[prod.0 as usize].push(nid);
346                    in_degree[nid.0 as usize] += 1;
347                }
348            }
349        }
350
351        // Min-heap on raw id for deterministic ordering.
352        let mut ready: BinaryHeap<std::cmp::Reverse<u32>> = in_degree
353            .iter()
354            .enumerate()
355            .filter(|(_, degree)| **degree == 0)
356            .map(|(raw, _)| std::cmp::Reverse(raw as u32))
357            .collect();
358
359        let mut order = Vec::with_capacity(self.nodes.len());
360        while let Some(std::cmp::Reverse(raw)) = ready.pop() {
361            let nid = NodeId(raw);
362            order.push(nid);
363            for &successor in &adj[raw as usize] {
364                let degree = &mut in_degree[successor.0 as usize];
365                *degree -= 1;
366                if *degree == 0 {
367                    ready.push(std::cmp::Reverse(successor.0));
368                }
369            }
370        }
371
372        if order.len() != self.nodes.len() {
373            return Err(GraphError::CycleDetected);
374        }
375        Ok(order)
376    }
377
378    // === Mutation API ===
379
380    /// Insert a node, wiring its producer/consumer edges. The node's `id`
381    /// field is overwritten with the freshly allocated [`NodeId`].
382    pub fn insert_node(&mut self, node: Node) -> NodeId {
383        let id = self.nodes.insert_with(|nid| {
384            let mut node = node;
385            node.id = nid;
386            node
387        });
388        self.connect_edges(id);
389        id
390    }
391
392    /// Remove a node, disconnecting its edges. Output values left with no
393    /// consumers (and not graph I/O or initializers) are deleted.
394    pub fn remove_node(&mut self, id: NodeId) {
395        if !self.nodes.contains(id) {
396            return;
397        }
398        self.disconnect_edges(id);
399        let outputs = self.node(id).outputs.clone();
400        self.nodes.remove(id);
401        for v in outputs {
402            self.gc_value_if_orphan(v);
403        }
404    }
405
406    /// Remove nodes in slice order.
407    ///
408    /// Each input edge is removed directly by `(NodeId, input_index)`, so this
409    /// remains linear in the number of removed edges even for a high-fanout
410    /// shared value.
411    pub fn remove_nodes(&mut self, ids: &[NodeId]) {
412        for &id in ids {
413            self.remove_node(id);
414        }
415    }
416
417    /// Replace disjoint node groups with one node each while updating shared
418    /// producer/consumer metadata in a batch.
419    ///
420    /// Each group is semantically equivalent to calling [`Graph::remove_node`]
421    /// for its IDs in slice order and then [`Graph::insert_node`] for the
422    /// replacement. In particular, replacement IDs and orphan-value collection
423    /// match that sequential operation. `graph_outputs` is retained for API
424    /// compatibility and checked against the per-value membership invariant in
425    /// debug builds.
426    pub fn replace_node_groups(
427        &mut self,
428        groups: Vec<(Vec<NodeId>, Node)>,
429        graph_outputs: &HashSet<ValueId>,
430    ) -> Vec<NodeId> {
431        debug_assert_eq!(
432            graph_outputs,
433            &self.outputs.iter().copied().collect::<HashSet<_>>()
434        );
435        let mut removed_nodes = HashSet::new();
436        for (node_ids, _) in &groups {
437            assert!(
438                !node_ids.is_empty(),
439                "replace_node_groups: group must not be empty"
440            );
441            for &id in node_ids {
442                assert!(
443                    self.nodes.contains(id),
444                    "replace_node_groups: node id not live"
445                );
446                assert!(
447                    removed_nodes.insert(id),
448                    "replace_node_groups: groups must be disjoint"
449                );
450            }
451        }
452
453        let mut inserted = Vec::with_capacity(groups.len());
454        for (node_ids, replacement) in groups {
455            for id in node_ids {
456                self.remove_node(id);
457            }
458            inserted.push(self.insert_node(replacement));
459        }
460
461        inserted
462    }
463
464    /// Replace node `old` in place with `new`, preserving the [`NodeId`].
465    ///
466    /// The old node's edges are disconnected and the new node's edges are
467    /// connected. Values that were outputs of `old` but not of `new` are left
468    /// in place (producer cleared); the caller may prune them.
469    pub fn replace_node(&mut self, old: NodeId, new: Node) -> NodeId {
470        assert!(self.nodes.contains(old), "replace_node: old id not live");
471        self.disconnect_edges(old);
472        {
473            let slot = self.nodes.get_mut(old).expect("old live");
474            let mut new = new;
475            new.id = old;
476            *slot = new;
477        }
478        self.connect_edges(old);
479        old
480    }
481
482    /// Splice `new_node` onto the edge feeding out of `value`:
483    /// `producer(value) → [new_node] → consumers(value)`.
484    ///
485    /// `new_node`'s single input becomes `value`, and it produces a fresh value
486    /// that replaces `value` in all of `value`'s original consumers.
487    pub fn insert_on_edge(&mut self, value: ValueId, new_node: Node) -> NodeId {
488        let (dtype, shape) = {
489            let v = self.value(value);
490            (v.dtype, v.shape.clone())
491        };
492        let new_value = self.create_value(dtype, shape);
493        // Redirect existing consumers onto the new value first (before the new
494        // node itself becomes a consumer of `value`).
495        self.replace_all_uses(value, new_value);
496
497        let mut new_node = new_node;
498        new_node.inputs = vec![Some(value)];
499        new_node.outputs = vec![new_value];
500        self.insert_node(new_node)
501    }
502
503    /// Replace one node input and update both values' consumer sets.
504    ///
505    /// This is constant-time on average for edge metadata. `None` disconnects
506    /// the slot and is used by node removal.
507    pub fn replace_input(
508        &mut self,
509        node: NodeId,
510        input_index: usize,
511        new_value: Option<ValueId>,
512    ) -> Option<ValueId> {
513        assert!(self.nodes.contains(node), "replace_input: node id not live");
514        assert!(
515            input_index < self.node(node).inputs.len(),
516            "replace_input: input index out of bounds"
517        );
518        if let Some(value) = new_value {
519            assert!(
520                self.values.contains(value),
521                "replace_input: value id not live"
522            );
523        }
524
525        let old_value = self.node(node).inputs[input_index];
526        if old_value == new_value {
527            return old_value;
528        }
529        if let Some(value) = old_value {
530            let removed = self
531                .value_mut(value)
532                .consumers
533                .remove(node, input_index as u32);
534            debug_assert!(removed, "old input edge must be present");
535        }
536        self.node_mut(node).inputs[input_index] = new_value;
537        if let Some(value) = new_value {
538            self.value_mut(value)
539                .consumers
540                .insert(node, input_index as u32);
541        }
542        old_value
543    }
544
545    /// Replace every use of `old_value` with `new_value` in consumer nodes and
546    /// in the graph output list, moving consumer edges accordingly.
547    pub fn replace_all_uses(&mut self, old_value: ValueId, new_value: ValueId) {
548        if old_value == new_value {
549            return;
550        }
551        let uses = match self.values.get(old_value) {
552            Some(value) => value.consumers.uses(),
553            None => return,
554        };
555        for (node, input_index) in uses {
556            if self.nodes.contains(node) {
557                self.replace_input(node, input_index as usize, Some(new_value));
558            }
559        }
560        if self.value(old_value).is_graph_output {
561            let mut outputs = self.outputs.clone();
562            for output in &mut outputs {
563                if *output == old_value {
564                    *output = new_value;
565                }
566            }
567            self.set_outputs(outputs);
568        }
569    }
570
571    // === Validation ===
572
573    /// Verify structural invariants (§3.3). Returns every defect found.
574    pub fn validate(&self) -> Result<(), Vec<GraphError>> {
575        let mut errors = Vec::new();
576        let graph_inputs: HashSet<_> = self.inputs.iter().copied().collect();
577        let graph_outputs: HashSet<_> = self.outputs.iter().copied().collect();
578        for (value, metadata) in self.values.iter() {
579            debug_assert_eq!(
580                metadata.is_graph_input,
581                graph_inputs.contains(&value),
582                "graph-input membership flag drifted for {value:?}"
583            );
584            debug_assert_eq!(
585                metadata.is_graph_output,
586                graph_outputs.contains(&value),
587                "graph-output membership flag drifted for {value:?}"
588            );
589        }
590
591        // 1. Node edges reference live values; collect produced values.
592        let mut produced: HashMap<ValueId, NodeId> = HashMap::new();
593        for (nid, node) in self.nodes.iter() {
594            for (input_index, input) in node.inputs.iter().enumerate() {
595                if let Some(value) = input {
596                    if !self.values.contains(*value) {
597                        errors.push(GraphError::DanglingValue(*value));
598                    } else if !self
599                        .value(*value)
600                        .consumers
601                        .contains(nid, input_index as u32)
602                    {
603                        errors.push(GraphError::ConsumerLinkMismatch(*value));
604                    }
605                }
606            }
607            for &v in &node.outputs {
608                if !self.values.contains(v) {
609                    errors.push(GraphError::DanglingValue(v));
610                    continue;
611                }
612                if produced.insert(v, nid).is_some() {
613                    errors.push(GraphError::DuplicateOutput(v));
614                }
615            }
616        }
617
618        // 2/3. Producer/consumer link consistency.
619        for (vid, val) in self.values.iter() {
620            if let Some(p) = val.producer {
621                if !self.nodes.contains(p) {
622                    errors.push(GraphError::DanglingNode(p));
623                } else if !self.node(p).outputs.contains(&vid) {
624                    errors.push(GraphError::ProducerLinkMismatch(vid));
625                }
626            }
627            for (consumer, input_index) in val.consumers.uses() {
628                if !self.nodes.contains(consumer) {
629                    errors.push(GraphError::DanglingNode(consumer));
630                } else if self.node(consumer).inputs.get(input_index as usize) != Some(&Some(vid)) {
631                    errors.push(GraphError::ConsumerLinkMismatch(vid));
632                }
633            }
634        }
635
636        // 4. Graph inputs must be sources.
637        for &inp in &self.inputs {
638            if let Some(val) = self.values.get(inp) {
639                if val.producer.is_some() {
640                    errors.push(GraphError::InputHasProducer(inp));
641                }
642                debug_assert!(val.is_graph_input);
643            } else {
644                errors.push(GraphError::DanglingValue(inp));
645            }
646        }
647
648        // 5. Graph outputs must be produced (unless they are graph inputs or
649        //    initializers passed straight through).
650        for &out in &self.outputs {
651            match self.values.get(out) {
652                Some(val) => {
653                    debug_assert!(val.is_graph_output);
654                    let is_source = val.is_graph_input || self.initializers.contains_key(&out);
655                    if val.producer.is_none() && !is_source {
656                        errors.push(GraphError::MissingProducer(out));
657                    }
658                }
659                None => errors.push(GraphError::DanglingValue(out)),
660            }
661        }
662
663        // 6. No cycles.
664        if let Err(e) = self.topological_order() {
665            errors.push(e);
666        }
667
668        // 7. Opset imports must have non-zero versions.
669        for (domain, &version) in &self.opset_imports {
670            if version == 0 {
671                errors.push(GraphError::InvalidOpsetImport {
672                    domain: domain.clone(),
673                    version,
674                });
675            }
676        }
677
678        // 8. Subgraphs validate recursively.
679        for sub in self.subgraphs.values() {
680            if let Err(mut sub_errors) = sub.validate() {
681                errors.append(&mut sub_errors);
682            }
683        }
684
685        if errors.is_empty() {
686            Ok(())
687        } else {
688            Err(errors)
689        }
690    }
691
692    // === Private edge maintenance ===
693
694    /// Wire a live node's edges into its input/output values.
695    fn connect_edges(&mut self, id: NodeId) {
696        let inputs = self.node(id).inputs.clone();
697        let outputs = self.node(id).outputs.clone();
698        for (input_index, input) in inputs.into_iter().enumerate() {
699            if let Some(value) = input
700                && let Some(metadata) = self.values.get_mut(value)
701            {
702                metadata.consumers.insert(id, input_index as u32);
703            }
704        }
705        for v in outputs {
706            if let Some(val) = self.values.get_mut(v) {
707                val.producer = Some(id);
708            }
709        }
710    }
711
712    /// Remove a live node's edges from its input/output values, without
713    /// deleting the values or the node itself.
714    fn disconnect_edges(&mut self, id: NodeId) {
715        let input_count = self.node(id).inputs.len();
716        let outputs = self.node(id).outputs.clone();
717        for input_index in 0..input_count {
718            self.replace_input(id, input_index, None);
719        }
720        for v in outputs {
721            if let Some(val) = self.values.get_mut(v)
722                && val.producer == Some(id)
723            {
724                val.producer = None;
725            }
726        }
727    }
728
729    /// Delete `value` if it has no producer, no consumers, and is not part of
730    /// the graph's I/O or initializers.
731    fn gc_value_if_orphan(&mut self, value: ValueId) {
732        let orphan = match self.values.get(value) {
733            Some(v) => {
734                v.producer.is_none()
735                    && v.consumers.is_empty()
736                    && !v.is_graph_input
737                    && !v.is_graph_output
738                    && !self.initializers.contains_key(&value)
739            }
740            None => false,
741        };
742        if orphan {
743            self.values.remove(value);
744            self.unknown_value_types.remove(&value);
745            self.unknown_value_shapes.remove(&value);
746        }
747    }
748}
749
750#[cfg(test)]
751mod tests {
752    use super::*;
753    use crate::shape::static_shape;
754    use crate::tensor::TensorData;
755
756    fn assert_graphs_identical(
757        mut batched: Graph,
758        mut sequential: Graph,
759        node_probes: usize,
760        value_probes: usize,
761        trial: usize,
762    ) {
763        assert_eq!(
764            format!("{batched:#?}"),
765            format!("{sequential:#?}"),
766            "graph mismatch on randomized trial {trial}"
767        );
768
769        // Arena free-list order is observable through subsequently allocated
770        // IDs, so exhaust the recycled slots as part of the equivalence check.
771        for _ in 0..node_probes {
772            let batched_id =
773                batched.insert_node(Node::new(NodeId(0), "Probe", Vec::new(), Vec::new()));
774            let sequential_id =
775                sequential.insert_node(Node::new(NodeId(0), "Probe", Vec::new(), Vec::new()));
776            assert_eq!(
777                batched_id, sequential_id,
778                "node arena mismatch on randomized trial {trial}"
779            );
780        }
781        for _ in 0..value_probes {
782            let batched_id = batched.create_value(DataType::Float32, static_shape([1]));
783            let sequential_id = sequential.create_value(DataType::Float32, static_shape([1]));
784            assert_eq!(
785                batched_id, sequential_id,
786                "value arena mismatch on randomized trial {trial}"
787            );
788        }
789    }
790
791    struct TestRng(u64);
792
793    impl TestRng {
794        fn next(&mut self) -> u64 {
795            self.0 ^= self.0 << 13;
796            self.0 ^= self.0 >> 7;
797            self.0 ^= self.0 << 17;
798            self.0
799        }
800
801        fn usize(&mut self, upper: usize) -> usize {
802            (self.next() as usize) % upper
803        }
804    }
805
806    fn reference_remove_node(graph: &mut Graph, id: NodeId) {
807        if !graph.nodes.contains(id) {
808            return;
809        }
810        let (inputs, outputs) = {
811            let node = graph.node(id);
812            (
813                node.input_values().collect::<Vec<_>>(),
814                node.outputs.clone(),
815            )
816        };
817        let unique_inputs: HashSet<_> = inputs.into_iter().collect();
818        for value in unique_inputs {
819            if let Some(metadata) = graph.values.get_mut(value) {
820                for (consumer, input_index) in metadata.consumers.uses() {
821                    if consumer == id {
822                        metadata.consumers.remove(consumer, input_index);
823                    }
824                }
825            }
826        }
827        for &value in &outputs {
828            if let Some(metadata) = graph.values.get_mut(value)
829                && metadata.producer == Some(id)
830            {
831                metadata.producer = None;
832            }
833        }
834        graph.nodes.remove(id);
835        for value in outputs {
836            let orphan = graph.values.get(value).is_some_and(|metadata| {
837                metadata.producer.is_none()
838                    && metadata.consumers.is_empty()
839                    && !graph.inputs.contains(&value)
840                    && !graph.outputs.contains(&value)
841                    && !graph.initializers.contains_key(&value)
842            });
843            if orphan {
844                graph.values.remove(value);
845                graph.unknown_value_types.remove(&value);
846                graph.unknown_value_shapes.remove(&value);
847            }
848        }
849    }
850
851    fn reference_topological_order(graph: &Graph) -> Result<Vec<NodeId>, GraphError> {
852        let mut in_degree: HashMap<NodeId, usize> =
853            graph.nodes.keys().map(|node| (node, 0)).collect();
854        let mut adjacency: HashMap<NodeId, Vec<NodeId>> = HashMap::new();
855        for (node_id, node) in graph.nodes.iter() {
856            for value in node.input_values() {
857                if let Some(producer) = graph.value(value).producer
858                    && graph.nodes.contains(producer)
859                {
860                    adjacency.entry(producer).or_default().push(node_id);
861                    *in_degree.get_mut(&node_id).unwrap() += 1;
862                }
863            }
864        }
865        let mut ready: BinaryHeap<std::cmp::Reverse<u32>> = in_degree
866            .iter()
867            .filter(|(_, degree)| **degree == 0)
868            .map(|(node, _)| std::cmp::Reverse(node.0))
869            .collect();
870        let mut order = Vec::with_capacity(graph.num_nodes());
871        while let Some(std::cmp::Reverse(raw)) = ready.pop() {
872            let node = NodeId(raw);
873            order.push(node);
874            if let Some(successors) = adjacency.get(&node) {
875                for successor in successors {
876                    let degree = in_degree.get_mut(successor).unwrap();
877                    *degree -= 1;
878                    if *degree == 0 {
879                        ready.push(std::cmp::Reverse(successor.0));
880                    }
881                }
882            }
883        }
884        if order.len() == graph.num_nodes() {
885            Ok(order)
886        } else {
887            Err(GraphError::CycleDetected)
888        }
889    }
890
891    /// Build `a -> Relu -> b -> Add(b, c) -> d`, returning ids.
892    fn sample_graph() -> Graph {
893        let mut g = Graph::new();
894        g.opset_imports.insert(String::new(), 17);
895        let a = g.create_named_value("a", DataType::Float32, static_shape([4]));
896        let c = g.create_named_value("c", DataType::Float32, static_shape([4]));
897        g.add_input(a);
898        g.add_input(c);
899
900        let b = g.create_value(DataType::Float32, static_shape([4]));
901        let relu = Node::new(NodeId(0), "Relu", vec![Some(a)], vec![b]);
902        g.insert_node(relu);
903
904        let d = g.create_named_value("d", DataType::Float32, static_shape([4]));
905        let add = Node::new(NodeId(0), "Add", vec![Some(b), Some(c)], vec![d]);
906        g.insert_node(add);
907        g.add_output(d);
908        g
909    }
910
911    #[test]
912    fn edges_are_wired_on_insert() {
913        let g = sample_graph();
914        assert_eq!(g.num_nodes(), 2);
915        // b has a producer (Relu) and a consumer (Add)
916        let b = g.value(ValueId(2));
917        assert!(b.producer.is_some());
918        assert_eq!(b.consumers.len(), 1);
919    }
920
921    #[test]
922    fn topo_order_is_valid_and_deterministic() {
923        let g = sample_graph();
924        let order = g.topological_order().unwrap();
925        assert_eq!(order.len(), 2);
926        // Relu (NodeId 0) must come before Add (NodeId 1)
927        assert_eq!(order, vec![NodeId(0), NodeId(1)]);
928    }
929
930    #[test]
931    fn uses_preserve_input_multiplicity_and_replace_one_slot() {
932        let mut graph = Graph::new();
933        let old = graph.create_value(DataType::Float32, static_shape([1]));
934        let new = graph.create_value(DataType::Float32, static_shape([1]));
935        graph.add_input(old);
936        graph.add_input(new);
937        let output = graph.create_value(DataType::Float32, static_shape([1]));
938        let node = graph.insert_node(Node::new(
939            NodeId(0),
940            "Add",
941            vec![Some(old), Some(old)],
942            vec![output],
943        ));
944
945        assert_eq!(graph.uses(old), vec![(node, 0), (node, 1)]);
946        assert_eq!(graph.consumers(old), vec![node]);
947        assert_eq!(graph.replace_input(node, 1, Some(new)), Some(old));
948        assert_eq!(graph.uses(old), vec![(node, 0)]);
949        assert_eq!(graph.uses(new), vec![(node, 1)]);
950        assert!(graph.validate().is_ok());
951
952        graph.remove_node(node);
953        assert!(graph.uses(old).is_empty());
954        assert!(graph.uses(new).is_empty());
955    }
956
957    #[test]
958    fn io_membership_flags_follow_ordered_lists() {
959        let mut graph = Graph::new();
960        let a = graph.create_value(DataType::Float32, static_shape([1]));
961        let b = graph.create_value(DataType::Float32, static_shape([1]));
962        graph.add_input(a);
963        graph.add_output(a);
964        graph.insert_output(0, b);
965        assert!(graph.value(a).is_graph_input);
966        assert!(graph.value(a).is_graph_output);
967        assert!(graph.value(b).is_graph_output);
968
969        assert_eq!(graph.remove_output(1), a);
970        assert!(!graph.value(a).is_graph_output);
971        graph.set_inputs(vec![b]);
972        assert!(!graph.value(a).is_graph_input);
973        assert!(graph.value(b).is_graph_input);
974        graph.set_outputs(vec![a]);
975        assert!(graph.value(a).is_graph_output);
976        assert!(!graph.value(b).is_graph_output);
977    }
978
979    #[test]
980    fn consumer_hash_insertion_order_is_not_observable() {
981        let mut first = Graph::new();
982        let hub = first.create_value(DataType::Float32, static_shape([1]));
983        first.add_input(hub);
984        let mut nodes = Vec::new();
985        for _ in 0..32 {
986            let output = first.create_value(DataType::Float32, static_shape([1]));
987            nodes.push(first.insert_node(Node::new(
988                NodeId(0),
989                "Add",
990                vec![Some(hub), Some(hub)],
991                vec![output],
992            )));
993        }
994        let mut shuffled = first.clone();
995        let mut uses = shuffled.uses(hub);
996        for &(node, input_index) in &uses {
997            shuffled.replace_input(node, input_index as usize, None);
998        }
999        let mut rng = TestRng(0x6a09_e667_f3bc_c909);
1000        for index in (1..uses.len()).rev() {
1001            uses.swap(index, rng.usize(index + 1));
1002        }
1003        for (node, input_index) in uses {
1004            shuffled.replace_input(node, input_index as usize, Some(hub));
1005        }
1006
1007        assert_eq!(first.uses(hub), shuffled.uses(hub));
1008        assert_eq!(first.consumers(hub), shuffled.consumers(hub));
1009        assert_eq!(first.topological_order(), shuffled.topological_order());
1010        assert_eq!(format!("{first:#?}"), format!("{shuffled:#?}"));
1011        assert_eq!(
1012            format!("{:#?}", sample_graph()),
1013            format!("{:#?}", sample_graph())
1014        );
1015        assert_eq!(nodes, first.consumers(hub));
1016    }
1017
1018    #[test]
1019    fn validate_accepts_wellformed_graph() {
1020        let g = sample_graph();
1021        assert!(g.validate().is_ok());
1022    }
1023
1024    #[test]
1025    fn predecessors_and_successors() {
1026        let g = sample_graph();
1027        assert_eq!(g.successors(NodeId(0)), vec![NodeId(1)]);
1028        assert_eq!(g.predecessors(NodeId(1)), vec![NodeId(0)]);
1029    }
1030
1031    #[test]
1032    fn nodes_between_walks_back() {
1033        let g = sample_graph();
1034        let between = g.nodes_between(&[ValueId(0), ValueId(1)], &[ValueId(3)]);
1035        assert!(between.contains(&NodeId(0)));
1036        assert!(between.contains(&NodeId(1)));
1037        assert_eq!(between.len(), 2);
1038    }
1039
1040    #[test]
1041    fn replace_all_uses_redirects_consumers() {
1042        let mut g = sample_graph();
1043        // New constant value replaces `b` as Add's input.
1044        let e = g.create_value(DataType::Float32, static_shape([4]));
1045        g.replace_all_uses(ValueId(2), e);
1046        // Add now consumes `e`, not `b`.
1047        let add = g.node(NodeId(1));
1048        assert!(add.input_values().any(|v| v == e));
1049        assert!(!add.input_values().any(|v| v == ValueId(2)));
1050        assert_eq!(g.consumers(e), vec![NodeId(1)]);
1051        assert!(g.value(ValueId(2)).consumers.is_empty());
1052    }
1053
1054    #[test]
1055    fn insert_on_edge_splices_node() {
1056        let mut g = sample_graph();
1057        // Insert an Identity between b (ValueId 2) and its consumer Add.
1058        let ident = Node::new(NodeId(0), "Identity", vec![], vec![]);
1059        let nid = g.insert_on_edge(ValueId(2), ident);
1060        // Add now consumes the new value produced by Identity.
1061        let new_out = g.node(nid).outputs[0];
1062        assert_eq!(g.value(new_out).producer, Some(nid));
1063        assert!(g.node(NodeId(1)).input_values().any(|v| v == new_out));
1064        assert!(g.validate().is_ok());
1065    }
1066
1067    #[test]
1068    fn remove_node_disconnects_and_gcs() {
1069        let mut g = sample_graph();
1070        g.remove_node(NodeId(1)); // remove Add
1071        assert_eq!(g.num_nodes(), 1);
1072        // `d` was Add's only output and a graph output -> kept, producer cleared
1073        assert!(g.value(ValueId(3)).producer.is_none());
1074        // b lost its consumer
1075        assert!(g.value(ValueId(2)).consumers.is_empty());
1076    }
1077
1078    #[test]
1079    fn remove_nodes_filters_wide_shared_input_once() {
1080        let mut graph = Graph::new();
1081        let input = graph.create_value(DataType::Float32, static_shape([1]));
1082        graph.add_input(input);
1083        let mut dead = Vec::new();
1084        let mut outputs = Vec::new();
1085        for _ in 0..1_000 {
1086            let output = graph.create_value(DataType::Float32, static_shape([1]));
1087            dead.push(graph.insert_node(Node::new(
1088                NodeId(0),
1089                "Relu",
1090                vec![Some(input)],
1091                vec![output],
1092            )));
1093            outputs.push(output);
1094        }
1095
1096        graph.remove_nodes(&dead);
1097
1098        assert_eq!(graph.num_nodes(), 0);
1099        assert!(graph.value(input).consumers.is_empty());
1100        assert!(
1101            outputs
1102                .into_iter()
1103                .all(|output| graph.try_value(output).is_none())
1104        );
1105    }
1106
1107    #[test]
1108    fn remove_nodes_keeps_surviving_shared_consumer() {
1109        let mut graph = Graph::new();
1110        let input = graph.create_value(DataType::Float32, static_shape([1]));
1111        graph.add_input(input);
1112        let mut nodes = Vec::new();
1113        for op_type in ["Relu", "Neg", "Abs"] {
1114            let output = graph.create_value(DataType::Float32, static_shape([1]));
1115            nodes.push(graph.insert_node(Node::new(
1116                NodeId(0),
1117                op_type,
1118                vec![Some(input)],
1119                vec![output],
1120            )));
1121        }
1122
1123        graph.remove_nodes(&nodes[..2]);
1124
1125        assert_eq!(graph.consumers(input), vec![nodes[2]]);
1126        assert!(graph.try_node(nodes[2]).is_some());
1127        assert!(graph.validate().is_ok());
1128    }
1129
1130    #[test]
1131    fn remove_nodes_collects_value_after_all_consumers_are_removed() {
1132        let mut graph = Graph::new();
1133        let input = graph.create_value(DataType::Float32, static_shape([1]));
1134        graph.add_input(input);
1135        let shared = graph.create_value(DataType::Float32, static_shape([1]));
1136        let producer = graph.insert_node(Node::new(
1137            NodeId(0),
1138            "Relu",
1139            vec![Some(input)],
1140            vec![shared],
1141        ));
1142        let mut consumers = Vec::new();
1143        for op_type in ["Neg", "Abs"] {
1144            let output = graph.create_value(DataType::Float32, static_shape([1]));
1145            consumers.push(graph.insert_node(Node::new(
1146                NodeId(0),
1147                op_type,
1148                vec![Some(shared)],
1149                vec![output],
1150            )));
1151        }
1152
1153        graph.remove_nodes(&[consumers[0], consumers[1], producer]);
1154
1155        assert!(graph.try_value(shared).is_none());
1156    }
1157
1158    #[test]
1159    fn remove_nodes_keeps_graph_outputs_and_initializers() {
1160        let mut graph = Graph::new();
1161        let input = graph.create_value(DataType::Float32, static_shape([1]));
1162        graph.add_input(input);
1163
1164        let graph_output = graph.create_value(DataType::Float32, static_shape([1]));
1165        let output_node = graph.insert_node(Node::new(
1166            NodeId(0),
1167            "Relu",
1168            vec![Some(input)],
1169            vec![graph_output],
1170        ));
1171        graph.add_output(graph_output);
1172
1173        let initializer = graph.create_value(DataType::Float32, static_shape([1]));
1174        let initializer_node = graph.insert_node(Node::new(
1175            NodeId(0),
1176            "Neg",
1177            vec![Some(input)],
1178            vec![initializer],
1179        ));
1180        graph.set_initializer(
1181            initializer,
1182            WeightRef::Inline(TensorData::from_raw(
1183                DataType::Float32,
1184                vec![1],
1185                0.0f32.to_le_bytes().to_vec(),
1186            )),
1187        );
1188
1189        graph.remove_nodes(&[output_node, initializer_node]);
1190
1191        assert!(graph.try_value(graph_output).is_some());
1192        assert!(graph.value(graph_output).producer.is_none());
1193        assert!(graph.try_value(initializer).is_some());
1194        assert!(graph.value(initializer).producer.is_none());
1195    }
1196
1197    #[test]
1198    fn remove_nodes_ignores_duplicate_and_nonlive_ids() {
1199        let graph = sample_graph();
1200        let ids = [NodeId(u32::MAX), NodeId(1), NodeId(1)];
1201        let mut sequential = graph.clone();
1202        for id in ids {
1203            sequential.remove_node(id);
1204        }
1205        let mut batched = graph;
1206        batched.remove_nodes(&ids);
1207
1208        assert_graphs_identical(batched, sequential, 3, 5, 0);
1209    }
1210
1211    #[test]
1212    fn remove_nodes_matches_sequential_removal_on_random_dags() {
1213        let mut rng = TestRng(0x4d59_5df4_d0f3_3173);
1214
1215        for trial in 0..10_000 {
1216            let input_count = 1 + rng.usize(3);
1217            let node_count = 1 + rng.usize(12);
1218            let mut graph = Graph::new();
1219            let mut values = Vec::new();
1220            for _ in 0..input_count {
1221                let input = graph.create_value(DataType::Float32, static_shape([1]));
1222                graph.add_input(input);
1223                values.push(input);
1224            }
1225
1226            let mut nodes = Vec::with_capacity(node_count);
1227            for _ in 0..node_count {
1228                let input_arity = 1 + rng.usize(3);
1229                let inputs = (0..input_arity)
1230                    .map(|_| Some(values[rng.usize(values.len())]))
1231                    .collect();
1232                let output = graph.create_value(DataType::Float32, static_shape([1]));
1233                nodes.push(graph.insert_node(Node::new(NodeId(0), "Random", inputs, vec![output])));
1234                if rng.usize(5) == 0 {
1235                    graph.mark_value_type_unknown(output);
1236                }
1237                if rng.usize(5) == 0 {
1238                    graph.mark_value_shape_unknown(output);
1239                }
1240                values.push(output);
1241            }
1242
1243            for _ in 0..rng.usize(4) {
1244                let output = values[rng.usize(values.len())];
1245                graph.add_output(output);
1246            }
1247
1248            for i in (1..nodes.len()).rev() {
1249                let j = rng.usize(i + 1);
1250                nodes.swap(i, j);
1251            }
1252            nodes.truncate(rng.usize(nodes.len() + 1));
1253
1254            let mut sequential = graph.clone();
1255            for &id in &nodes {
1256                sequential.remove_node(id);
1257            }
1258            let mut batched = graph;
1259            batched.remove_nodes(&nodes);
1260
1261            assert_graphs_identical(
1262                batched,
1263                sequential,
1264                node_count + 1,
1265                input_count + node_count + 1,
1266                trial,
1267            );
1268        }
1269    }
1270
1271    #[test]
1272    fn single_node_removal_matches_vector_reference_on_random_dags() {
1273        let mut rng = TestRng(0xbb67_ae85_84ca_a73b);
1274        for trial in 0..2_000 {
1275            let mut graph = Graph::new();
1276            let input_count = 1 + rng.usize(3);
1277            let node_count = 1 + rng.usize(24);
1278            let mut values = Vec::new();
1279            for _ in 0..input_count {
1280                let value = graph.create_value(DataType::Float32, static_shape([1]));
1281                graph.add_input(value);
1282                values.push(value);
1283            }
1284            let mut nodes = Vec::new();
1285            for _ in 0..node_count {
1286                let input_count = 1 + rng.usize(4);
1287                let inputs = (0..input_count)
1288                    .map(|_| Some(values[rng.usize(values.len())]))
1289                    .collect();
1290                let output = graph.create_value(DataType::Float32, static_shape([1]));
1291                nodes.push(graph.insert_node(Node::new(NodeId(0), "Random", inputs, vec![output])));
1292                values.push(output);
1293            }
1294            for _ in 0..rng.usize(4) {
1295                graph.add_output(values[rng.usize(values.len())]);
1296            }
1297
1298            let mut removals = nodes.clone();
1299            for index in (1..removals.len()).rev() {
1300                removals.swap(index, rng.usize(index + 1));
1301            }
1302            removals.truncate(rng.usize(removals.len() + 1));
1303            if let Some(&duplicate) = removals.first() {
1304                removals.push(duplicate);
1305            }
1306            removals.push(NodeId(u32::MAX));
1307
1308            let mut actual = graph.clone();
1309            let mut reference = graph;
1310            for &node in &removals {
1311                actual.remove_node(node);
1312                reference_remove_node(&mut reference, node);
1313            }
1314
1315            assert_eq!(
1316                format!("{actual:#?}"),
1317                format!("{reference:#?}"),
1318                "debug mismatch on trial {trial}"
1319            );
1320            assert_eq!(
1321                actual.topological_order(),
1322                reference.topological_order(),
1323                "topology mismatch on trial {trial}"
1324            );
1325            for value in actual.values.keys() {
1326                assert_eq!(
1327                    actual.uses(value),
1328                    reference.uses(value),
1329                    "uses mismatch for {value:?} on trial {trial}"
1330                );
1331                assert_eq!(
1332                    actual.consumers(value),
1333                    reference.consumers(value),
1334                    "consumers mismatch for {value:?} on trial {trial}"
1335                );
1336            }
1337        }
1338    }
1339
1340    #[test]
1341    fn vec_indexed_topology_matches_hashmap_reference_on_random_dags() {
1342        let mut rng = TestRng(0x3c6e_f372_fe94_f82b);
1343        for trial in 0..2_000 {
1344            let mut graph = Graph::new();
1345            let input = graph.create_value(DataType::Float32, static_shape([1]));
1346            graph.add_input(input);
1347            let mut values = vec![input];
1348            let mut nodes = Vec::new();
1349            for _ in 0..(1 + rng.usize(48)) {
1350                let inputs = (0..(1 + rng.usize(4)))
1351                    .map(|_| Some(values[rng.usize(values.len())]))
1352                    .collect();
1353                let output = graph.create_value(DataType::Float32, static_shape([1]));
1354                nodes.push(graph.insert_node(Node::new(NodeId(0), "Random", inputs, vec![output])));
1355                values.push(output);
1356            }
1357            for &node in nodes.iter().filter(|_| rng.usize(5) == 0) {
1358                graph.remove_node(node);
1359            }
1360            assert_eq!(
1361                graph.topological_order(),
1362                reference_topological_order(&graph),
1363                "topology mismatch on trial {trial}"
1364            );
1365        }
1366    }
1367
1368    #[test]
1369    fn replace_node_groups_matches_sequential_mutation() {
1370        let mut sequential = Graph::new();
1371        let input = sequential.create_value(DataType::Float32, static_shape([1]));
1372        sequential.add_input(input);
1373        let interior = sequential.create_value(DataType::Float32, static_shape([1]));
1374        let first = sequential.insert_node(Node::new(
1375            NodeId(0),
1376            "Relu",
1377            vec![Some(input)],
1378            vec![interior],
1379        ));
1380        let output = sequential.create_value(DataType::Float32, static_shape([1]));
1381        let second = sequential.insert_node(Node::new(
1382            NodeId(0),
1383            "Relu",
1384            vec![Some(interior)],
1385            vec![output],
1386        ));
1387        sequential.add_output(output);
1388        let sibling_output = sequential.create_value(DataType::Float32, static_shape([1]));
1389        let sibling = sequential.insert_node(Node::new(
1390            NodeId(0),
1391            "Neg",
1392            vec![Some(input)],
1393            vec![sibling_output],
1394        ));
1395        sequential.add_output(sibling_output);
1396
1397        let mut batched = sequential.clone();
1398        let graph_outputs: HashSet<_> = batched.outputs.iter().copied().collect();
1399        let replacement0 = Node::new(NodeId(0), "EPContext", vec![Some(input)], vec![output]);
1400        let replacement1 = Node::new(
1401            NodeId(0),
1402            "EPContext",
1403            vec![Some(input)],
1404            vec![sibling_output],
1405        );
1406
1407        sequential.remove_node(first);
1408        sequential.remove_node(second);
1409        let replacement0_id = sequential.insert_node(replacement0.clone());
1410        sequential.remove_node(sibling);
1411        let replacement1_id = sequential.insert_node(replacement1.clone());
1412
1413        let inserted = batched.replace_node_groups(
1414            vec![
1415                (vec![first, second], replacement0),
1416                (vec![sibling], replacement1),
1417            ],
1418            &graph_outputs,
1419        );
1420        assert_eq!(inserted, vec![replacement0_id, replacement1_id]);
1421
1422        let sequential_nodes: Vec<_> = sequential
1423            .nodes
1424            .iter()
1425            .map(|(id, node)| {
1426                (
1427                    id,
1428                    node.op_type.clone(),
1429                    node.inputs.clone(),
1430                    node.outputs.clone(),
1431                )
1432            })
1433            .collect();
1434        let batched_nodes: Vec<_> = batched
1435            .nodes
1436            .iter()
1437            .map(|(id, node)| {
1438                (
1439                    id,
1440                    node.op_type.clone(),
1441                    node.inputs.clone(),
1442                    node.outputs.clone(),
1443                )
1444            })
1445            .collect();
1446        assert_eq!(batched_nodes, sequential_nodes);
1447
1448        let sequential_values: Vec<_> = sequential
1449            .values
1450            .iter()
1451            .map(|(id, value)| (id, value.producer, value.consumers.clone()))
1452            .collect();
1453        let batched_values: Vec<_> = batched
1454            .values
1455            .iter()
1456            .map(|(id, value)| (id, value.producer, value.consumers.clone()))
1457            .collect();
1458        assert_eq!(batched_values, sequential_values);
1459
1460        let next_sequential =
1461            sequential.insert_node(Node::new(NodeId(0), "Identity", Vec::new(), Vec::new()));
1462        let next_batched =
1463            batched.insert_node(Node::new(NodeId(0), "Identity", Vec::new(), Vec::new()));
1464        assert_eq!(next_batched, next_sequential);
1465    }
1466
1467    #[test]
1468    fn replace_node_groups_matches_sequential_orphan_collection() {
1469        let mut sequential = Graph::new();
1470        let input = sequential.create_value(DataType::Float32, static_shape([1]));
1471        sequential.add_input(input);
1472        let interior = sequential.create_value(DataType::Float32, static_shape([1]));
1473        let producer = sequential.insert_node(Node::new(
1474            NodeId(0),
1475            "Relu",
1476            vec![Some(input)],
1477            vec![interior],
1478        ));
1479        let output = sequential.create_value(DataType::Float32, static_shape([1]));
1480        let consumer = sequential.insert_node(Node::new(
1481            NodeId(0),
1482            "Relu",
1483            vec![Some(interior)],
1484            vec![output],
1485        ));
1486        sequential.add_output(output);
1487
1488        let mut batched = sequential.clone();
1489        let graph_outputs: HashSet<_> = batched.outputs.iter().copied().collect();
1490        let replacement = Node::new(NodeId(0), "EPContext", vec![Some(input)], vec![output]);
1491
1492        sequential.remove_node(consumer);
1493        sequential.remove_node(producer);
1494        sequential.insert_node(replacement.clone());
1495        batched.replace_node_groups(
1496            vec![(vec![consumer, producer], replacement)],
1497            &graph_outputs,
1498        );
1499
1500        assert!(sequential.try_value(interior).is_none());
1501        assert!(batched.try_value(interior).is_none());
1502        let next_sequential = sequential.create_value(DataType::Float32, static_shape([1]));
1503        let next_batched = batched.create_value(DataType::Float32, static_shape([1]));
1504        assert_eq!(next_batched, next_sequential);
1505    }
1506
1507    #[test]
1508    fn replace_node_preserves_id() {
1509        let mut g = sample_graph();
1510        let d = g.node(NodeId(1)).outputs[0];
1511        let b = g.node(NodeId(1)).inputs[0];
1512        let c = g.node(NodeId(1)).inputs[1];
1513        let sub = Node::new(NodeId(0), "Sub", vec![b, c], vec![d]);
1514        let id = g.replace_node(NodeId(1), sub);
1515        assert_eq!(id, NodeId(1));
1516        assert_eq!(g.node(NodeId(1)).op_type, "Sub");
1517        assert!(g.validate().is_ok());
1518    }
1519
1520    #[test]
1521    fn cycle_is_detected() {
1522        let mut g = Graph::new();
1523        let v0 = g.create_value(DataType::Float32, static_shape([1]));
1524        let v1 = g.create_value(DataType::Float32, static_shape([1]));
1525        // n0: v1 -> v0 ; n1: v0 -> v1  (cycle)
1526        g.insert_node(Node::new(NodeId(0), "A", vec![Some(v1)], vec![v0]));
1527        g.insert_node(Node::new(NodeId(0), "B", vec![Some(v0)], vec![v1]));
1528        assert_eq!(g.topological_order(), Err(GraphError::CycleDetected));
1529        assert!(g.validate().is_err());
1530    }
1531
1532    #[test]
1533    fn intern_symbol_dedups_by_name() {
1534        let mut g = Graph::new();
1535        let s1 = g.intern_symbol("batch");
1536        let s2 = g.intern_symbol("batch");
1537        let s3 = g.intern_symbol("seq");
1538        assert_eq!(s1, s2);
1539        assert_ne!(s1, s3);
1540    }
1541}