Resolve scopes before flattening network (#4235)

* Compiler changes

* Run preprocessor on wrapping network

* Remove test nodes

* FIx

* move generate_node_paths call to compiler

* Make it more obvious what input is set to

node-graph/graph-craft/src/document.rs

@@ -655,6 +655,73 @@ impl NodeNetwork {
 	}
 }
 
+// Functions for resolving scope inputs
+impl NodeNetwork {
+	pub fn resolve_scope_inputs(&mut self) {
+		let mut leftover = Vec::new();
+		self.resolve_scope_inputs_impl(None, &mut leftover);
+		assert!(leftover.is_empty(), "Unresolved scope keys at top level: {leftover:?}");
+	}
+
+	fn resolve_scope_inputs_impl(&mut self, parent: Option<&ScopeChain<'_>>, network_inputs: &mut Vec<NodeInput>) {
+		let scope_chain = ScopeChain {
+			scopes: &self.scope_injections,
+			parent,
+		};
+
+		for node in self.nodes.values_mut() {
+			let DocumentNodeImplementation::Network(network) = &mut node.implementation else { continue };
+			network.resolve_scope_inputs_impl(Some(&scope_chain), &mut node.inputs);
+		}
+
+		let mut key_to_idx: FxHashMap<Cow<'static, str>, usize> = FxHashMap::default();
+		for node in self.nodes.values_mut() {
+			for input in node.inputs.iter_mut() {
+				let NodeInput::Scope(key) = input else { continue };
+				*input = match scope_chain.lookup(key) {
+					ScopeChainLookup::Current(node_id, _ty) => NodeInput::node(*node_id, 0),
+					ScopeChainLookup::Parent(_node_id, ty) => {
+						let import_index = *key_to_idx.entry(key.clone()).or_insert_with(|| {
+							let index = network_inputs.len();
+							network_inputs.push(NodeInput::Scope(key.clone()));
+							index
+						});
+						NodeInput::Import {
+							import_type: ty.clone(),
+							import_index,
+						}
+					}
+					ScopeChainLookup::None => panic!("Scope key `{key}` not found in any ancestor scope_injections"),
+				};
+			}
+		}
+	}
+}
+
+struct ScopeChain<'a> {
+	scopes: &'a FxHashMap<String, (NodeId, Type)>,
+	parent: Option<&'a ScopeChain<'a>>,
+}
+enum ScopeChainLookup<'a> {
+	Current(&'a NodeId, &'a Type),
+	Parent(&'a NodeId, &'a Type),
+	None,
+}
+impl ScopeChain<'_> {
+	fn lookup(&'_ self, key: &str) -> ScopeChainLookup<'_> {
+		self.scopes
+			.get(key)
+			.map(|(id, ty)| ScopeChainLookup::Current(id, ty))
+			.or_else(|| {
+				self.parent.and_then(|parent| match parent.lookup(key) {
+					ScopeChainLookup::Current(id, ty) | ScopeChainLookup::Parent(id, ty) => Some(ScopeChainLookup::Parent(id, ty)),
+					ScopeChainLookup::None => None,
+				})
+			})
+			.unwrap_or(ScopeChainLookup::None)
+	}
+}
+
 /// Functions for compiling the network
 impl NodeNetwork {
 	/// Replace all references in the graph of a node ID with a new node ID defined by the function `f`.
@@ -784,18 +851,6 @@ impl NodeNetwork {
 		are_inputs_used
 	}
 
-	pub fn resolve_scope_inputs(&mut self) {
-		for node in self.nodes.values_mut() {
-			for input in node.inputs.iter_mut() {
-				if let NodeInput::Scope(key) = input {
-					let (import_id, _ty) = self.scope_injections.get(key.as_ref()).expect("Tried to import a non existent key from scope");
-					// TODO use correct output index
-					*input = NodeInput::node(*import_id, 0);
-				}
-			}
-		}
-	}
-
 	/// Remove all nodes that contain [`DocumentNodeImplementation::Network`] by moving the nested nodes into the parent network.
 	pub fn flatten(&mut self, node_id: NodeId) {
 		self.flatten_with_fns(node_id, merge_ids, NodeId::new)
@@ -886,11 +941,7 @@ impl NodeNetwork {
 						}
 						NodeInput::Value { .. } => unreachable!("Value inputs should have been replaced with value nodes"),
 						NodeInput::Inline(_) => (),
-						NodeInput::Scope(ref key) => {
-							let (import_id, _ty) = self.scope_injections.get(key.as_ref()).expect("Tried to import a non existent key from scope");
-							// TODO use correct output index
-							nested_node.inputs[nested_input_index] = NodeInput::node(*import_id, 0);
-						}
+						NodeInput::Scope(_) => unreachable!("Scope inputs should have been resolved by resolve_scope_inputs_recursive before flattening"),
 						NodeInput::Reflection(_) => unreachable!("Reflection inputs should have been replaced with value nodes"),
 					}
 				}
@@ -906,21 +957,8 @@ impl NodeNetwork {
 		}
 	}
 
+	/// Connect all nodes that were previously connected to this node to the nodes of the inner network
 	fn replace_node_with_its_exports(&mut self, id: NodeId, original_location: &OriginalLocation, exports: &[NodeInput]) {
-		// Connect scope injections to the inner network export
-		self.scope_injections.values_mut().for_each(|(node_id, _ty)| {
-			if node_id == &id {
-				let Some(export) = exports.first() else {
-					log::error!("Inner network should have at least one export");
-					return;
-				};
-				if let NodeInput::Node { node_id: export_id, output_index: _ } = export {
-					*node_id = *export_id;
-				}
-			}
-		});
-
-		// Connect all nodes that were previously connected to this node to the nodes of the inner network
 		for (i, export) in exports.iter().enumerate() {
 			if let NodeInput::Node { node_id, output_index, .. } = &export {
 				for deps in &original_location.dependants {

node-graph/graph-craft/src/graphene_compiler.rs

@@ -6,12 +6,13 @@ pub struct Compiler {}
 
 impl Compiler {
 	pub fn compile(&self, mut network: NodeNetwork) -> impl Iterator<Item = Result<ProtoNetwork, String>> {
+		network.resolve_scope_inputs();
+		network.generate_node_paths(&[]);
 		let node_ids = network.nodes.keys().copied().collect::<Vec<_>>();
 		network.populate_dependants();
 		for id in node_ids {
 			network.flatten(id);
 		}
-		network.resolve_scope_inputs();
 		network.remove_redundant_passthrough_nodes();
 		// network.remove_dead_nodes(0);
 		let proto_networks = network.into_proto_networks();