Eliminate redundant graph edge computation in the compilation pipeline

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)
 		})
 	}

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<NodeId, Vec<NodeId>> {
-		let mut edges: HashMap<NodeId, Vec<NodeId>> = 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<Vec<NodeId>>) {
 		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<NodeId, Vec<NodeId>> {
-		let mut edges: HashMap<NodeId, Vec<NodeId>> = 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<Vec<usize>>, FxHashMap<NodeId, usize>) {
 		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<Vec<Vec<NodeId>>, 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<NodeId, Vec<NodeId>>, 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<NodeId>], 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<Vec<NodeId>> {
+		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<NodeId>, FxHashMap<NodeId, usize>), 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();