diff --git a/node-graph/graph-craft/src/graphene_compiler.rs b/node-graph/graph-craft/src/graphene_compiler.rs index 5041b34553..114fb8f691 100644 --- a/node-graph/graph-craft/src/graphene_compiler.rs +++ b/node-graph/graph-craft/src/graphene_compiler.rs @@ -17,8 +17,8 @@ impl Compiler { let proto_networks = network.into_proto_networks(); proto_networks.map(move |mut proto_network| { - proto_network.insert_context_nullification_nodes()?; - proto_network.generate_stable_node_ids(); + let outwards_edges = proto_network.insert_context_nullification_nodes()?; + proto_network.generate_stable_node_ids(&outwards_edges); Ok(proto_network) }) } diff --git a/node-graph/graph-craft/src/proto.rs b/node-graph/graph-craft/src/proto.rs index 7e3250590a..e6a16cf7f8 100644 --- a/node-graph/graph-craft/src/proto.rs +++ b/node-graph/graph-craft/src/proto.rs @@ -5,7 +5,7 @@ pub use core_types::registry::*; use core_types::*; use rustc_hash::FxHashMap; use std::borrow::Cow; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::fmt::Debug; use std::hash::Hash; @@ -232,49 +232,19 @@ impl ProtoNetwork { } /// Construct a hashmap containing a list of the nodes that depend on this proto network. - pub fn collect_outwards_edges(&self) -> HashMap> { - let mut edges: HashMap> = HashMap::new(); - for (id, node) in &self.nodes { - if let ConstructionArgs::Nodes(ref_nodes) = &node.construction_args { - for ref_id in ref_nodes { - self.check_ref(ref_id, id); - edges.entry(*ref_id).or_default().push(*id) - } - } - } - edges - } - /// Convert all node IDs to be stable (based on the hash generated by [`ProtoNode::stable_node_id`]). - /// This function requires that the graph be topologically sorted. - pub fn generate_stable_node_ids(&mut self) { - debug_assert!(self.is_topologically_sorted()); - let outwards_edges = self.collect_outwards_edges(); - + /// This function requires that the graph be topologically sorted with dense sequential IDs (0..N). + /// Accepts pre-computed outwards edges to avoid redundant graph traversal. + pub fn generate_stable_node_ids(&mut self, outwards_edges: &Vec>) { for index in 0..self.nodes.len() { let Some(sni) = self.nodes[index].1.stable_node_id() else { panic!("failed to generate stable node id for node {:#?}", self.nodes[index].1); }; - self.replace_node_id(&outwards_edges, NodeId(index as u64), sni); + self.replace_node_id(outwards_edges, NodeId(index as u64), sni); self.nodes[index].0 = sni; } } - // TODO: Remove - /// Create a hashmap with the list of nodes this proto network depends on/uses as inputs. - pub fn collect_inwards_edges(&self) -> HashMap> { - let mut edges: HashMap> = HashMap::new(); - for (id, node) in &self.nodes { - if let ConstructionArgs::Nodes(ref_nodes) = &node.construction_args { - for ref_id in ref_nodes { - self.check_ref(ref_id, id); - edges.entry(*id).or_default().push(*ref_id) - } - } - } - edges - } - fn collect_inwards_edges_with_mapping(&self) -> (Vec>, FxHashMap) { let id_map: FxHashMap<_, _> = self.nodes.iter().enumerate().map(|(idx, (id, _))| (*id, idx)).collect(); @@ -296,7 +266,8 @@ impl ProtoNetwork { /// Inserts context nullification nodes to optimize caching. /// This analysis is performed after topological sorting to ensure proper dependency tracking. - pub fn insert_context_nullification_nodes(&mut self) -> Result<(), String> { + /// Returns the outwards edges of the final sorted graph for reuse by subsequent passes. + pub fn insert_context_nullification_nodes(&mut self) -> Result>, String> { // Perform topological sort once self.reorder_ids()?; @@ -305,7 +276,8 @@ impl ProtoNetwork { // Perform topological sort a second time to integrate the new nodes self.reorder_ids()?; - Ok(()) + // Compute outwards edges on the now-sorted dense graph (node IDs are 0..N) + Ok(self.collect_outwards_edges_dense()) } fn insert_context_nullification_node(&mut self, node_id: NodeId, context_deps: ContextFeatures) -> NodeId { @@ -425,13 +397,11 @@ impl ProtoNetwork { } /// Update all of the references to a node ID in the graph with a new ID named `compose_node_id`. - fn replace_node_id(&mut self, outwards_edges: &HashMap>, node_id: NodeId, replacement_node_id: NodeId) { - // Update references in other nodes to use the new node - if let Some(referring_nodes) = outwards_edges.get(&node_id) { - for &referring_node_id in referring_nodes { - let (_, referring_node) = &mut self.nodes[referring_node_id.0 as usize]; - referring_node.map_ids(|id| if id == node_id { replacement_node_id } else { id }) - } + /// Uses dense Vec-indexed outwards edges (requires sequential node IDs 0..N). + fn replace_node_id(&mut self, outwards_edges: &[Vec], node_id: NodeId, replacement_node_id: NodeId) { + for &referring_node_id in &outwards_edges[node_id.0 as usize] { + let (_, referring_node) = &mut self.nodes[referring_node_id.0 as usize]; + referring_node.map_ids(|id| if id == node_id { replacement_node_id } else { id }) } if self.output == node_id { @@ -445,6 +415,21 @@ impl ProtoNetwork { }); } + /// Collect outwards edges using dense Vec indexing. Requires node IDs to be sequential 0..N + /// (i.e., the graph has been through `reorder_ids`). + fn collect_outwards_edges_dense(&self) -> Vec> { + let mut edges = vec![Vec::new(); self.nodes.len()]; + for (id, node) in &self.nodes { + if let ConstructionArgs::Nodes(ref_nodes) = &node.construction_args { + for ref_id in ref_nodes { + self.check_ref(ref_id, id); + edges[ref_id.0 as usize].push(*id); + } + } + } + edges + } + // Based on https://en.wikipedia.org/wiki/Topological_sorting#Depth-first_search // This approach excludes nodes that are not connected pub fn topological_sort(&self) -> Result<(Vec, FxHashMap), String> { @@ -483,24 +468,6 @@ impl ProtoNetwork { Ok((sorted, id_map)) } - fn is_topologically_sorted(&self) -> bool { - let mut visited = HashSet::new(); - - let inwards_edges = self.collect_inwards_edges(); - for (id, _) in &self.nodes { - for &dependency in inwards_edges.get(id).unwrap_or(&Vec::new()) { - if !visited.contains(&dependency) { - dbg!(id, dependency); - dbg!(&visited); - dbg!(&self.nodes); - return false; - } - } - visited.insert(*id); - } - true - } - /// Sort the nodes vec so it is in a topological order. This ensures that no node takes an input from a node that is found later in the list. fn reorder_ids(&mut self) -> Result<(), String> { let (order, _id_map) = self.topological_sort()?; @@ -941,10 +908,10 @@ mod test { #[test] fn stable_node_id_generation() { let mut construction_network = test_network(); - construction_network + let outwards_edges = construction_network .insert_context_nullification_nodes() .expect("Error when calling 'insert_context_nullification_nodes' on 'construction_network."); - construction_network.generate_stable_node_ids(); + construction_network.generate_stable_node_ids(&outwards_edges); assert_eq!(construction_network.nodes[0].1.identifier.as_str(), "value"); let ids: Vec<_> = construction_network.nodes.iter().map(|(id, _)| *id).collect();