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

editor/src/node_graph_executor/runtime.rs

@@ -345,13 +345,13 @@ impl NodeRuntime {
 		None
 	}
 
-	async fn update_network(&mut self, mut graph: NodeNetwork) -> Result<ResolvedDocumentNodeTypesDelta, (ResolvedDocumentNodeTypesDelta, String)> {
-		if let Err(e) = self.preprocessor.expand_network(&mut graph, &self.resources) {
+	async fn update_network(&mut self, graph: NodeNetwork) -> Result<ResolvedDocumentNodeTypesDelta, (ResolvedDocumentNodeTypesDelta, String)> {
+		let mut scoped_network = wrap_network_in_scope(graph, self.editor_api.clone());
+
+		if let Err(e) = self.preprocessor.preprocess(&mut scoped_network, &self.resources) {
 			return Err((ResolvedDocumentNodeTypesDelta::default(), e.to_string()));
 		}
 
-		let scoped_network = wrap_network_in_scope(graph, self.editor_api.clone());
-
 		// We assume only one output
 		assert_eq!(scoped_network.exports.len(), 1, "Graph with multiple outputs not yet handled");
 

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();

node-graph/graphene-cli/src/main.rs

@@ -242,10 +242,10 @@ fn compile_graph(document_string: String, editor_api: Arc<PlatformEditorApi>) ->
 	let mut network = load_network(&document_string);
 	fix_nodes(&mut network);
 
-	let preprocessor = preprocessor::Preprocessor::new();
-	preprocessor.expand_network(&mut network, &ResourceRegistry::default()).expect("Failed to expand network"); // TODO: actually load the resources from the document
+	let mut wrapped_network = wrap_network_in_scope(network, editor_api);
 
-	let wrapped_network = wrap_network_in_scope(network.clone(), editor_api);
+	let preprocessor = preprocessor::Preprocessor::new();
+	preprocessor.preprocess(&mut wrapped_network, &ResourceRegistry::default()).expect("Failed to expand network"); // TODO: actually load the resources from the document
 
 	let compiler = Compiler {};
 	compiler.compile_single(wrapped_network).map_err(|x| x.into())

node-graph/interpreted-executor/benches/benchmark_util.rs

@@ -9,11 +9,11 @@ use interpreted_executor::dynamic_executor::DynamicExecutor;
 use interpreted_executor::util::wrap_network_in_scope;
 
 pub fn setup_network(name: &str) -> (DynamicExecutor, ProtoNetwork) {
-	let mut network = load_from_name(name);
+	let network = load_from_name(name);
 	let editor_api = std::sync::Arc::new(EditorApi::default());
+	let mut network = wrap_network_in_scope(network, editor_api);
 	let preprocessor = preprocessor::Preprocessor::new();
-	preprocessor.expand_network(&mut network, &ResourceRegistry::default()).unwrap();
-	let network = wrap_network_in_scope(network, editor_api);
+	preprocessor.preprocess(&mut network, &ResourceRegistry::default()).unwrap();
 	let proto_network = compile(network);
 	let executor = block_on(DynamicExecutor::new(proto_network.clone())).unwrap();
 	(executor, proto_network)

node-graph/interpreted-executor/src/util.rs

@@ -10,9 +10,7 @@ use graphene_std::uuid::NodeId;
 use std::sync::Arc;
 use wgpu_executor::WgpuExecutor;
 
-pub fn wrap_network_in_scope(mut network: NodeNetwork, editor_api: Arc<PlatformEditorApi>) -> NodeNetwork {
-	network.generate_node_paths(&[]);
-
+pub fn wrap_network_in_scope(network: NodeNetwork, editor_api: Arc<PlatformEditorApi>) -> NodeNetwork {
 	let inner_network = DocumentNode {
 		implementation: DocumentNodeImplementation::Network(network),
 		inputs: vec![],
@@ -126,7 +124,6 @@ pub fn wrap_network_in_scope(mut network: NodeNetwork, editor_api: Arc<PlatformE
 		exports: vec![NodeInput::node(NodeId(1), 0)],
 		nodes: nodes.into_iter().enumerate().map(|(id, node)| (NodeId(id as u64), node)).collect(),
 		scope_injections: scope_injections.into_iter().collect(),
-		// TODO(TrueDoctor): check if it makes sense to set `generated` to `true`
-		generated: false,
+		generated: true,
 	}
 }

node-graph/nodes/gcore/src/debug.rs

@@ -34,13 +34,3 @@ fn unwrap_option<T: Default>(_: impl Ctx, #[implementations(Option<f64>, Option<
 fn clone<'i, T: Clone + 'i>(_: impl Ctx, #[implementations(&List<Raster<CPU>>)] value: &'i T) -> T {
 	value.clone()
 }
-
-#[node_macro::node(category("Debug"), inject_scope)]
-fn inject_scope_test_producer(_: impl Ctx) -> bool {
-	true
-}
-
-#[node_macro::node(category("Debug"))]
-fn inject_scope_test_consumer(_: impl Ctx, _primary: (), #[scope(inject_scope_test_producer::IDENTIFIER)] injected: bool) -> bool {
-	injected
-}

node-graph/preprocessor/src/lib.rs

@@ -19,22 +19,35 @@ pub struct Preprocessor {
 }
 
 impl Preprocessor {
-	pub fn expand_network(&self, network: &mut NodeNetwork, resources: &ResourceRegistry) -> Result<(), PreprocessorError> {
+	pub fn preprocess(&self, network: &mut NodeNetwork, resources: &ResourceRegistry) -> Result<(), PreprocessorError> {
 		self.insert_inject_scopes(network);
-		replace_resource_inputs(network, resources)?;
-		self.expand_network_inner(network);
+		self.replace_resource_inputs(network, resources)?;
+		self.expand_network(network);
 		Ok(())
 	}
 }
 
+impl Preprocessor {
+	fn insert_inject_scopes(&self, network: &mut NodeNetwork) {
+		for (identifier, (template, ty)) in self.inject_scopes.iter() {
+			let mut hasher = DefaultHasher::new();
+
+			identifier.as_str().hash(&mut hasher);
+			let producer_id = NodeId(hasher.finish());
+			network.nodes.insert(producer_id, template.clone());
+
+			network.scope_injections.insert(identifier.as_str().to_string(), (producer_id, ty.clone()));
+		}
+	}
+
 	/// Replace every `TaggedValue::Resource(hash)` input with a reference to a freshly inserted `resource` proto node.
-fn replace_resource_inputs(network: &mut NodeNetwork, resources: &ResourceRegistry) -> Result<(), PreprocessorError> {
+	fn replace_resource_inputs(&self, network: &mut NodeNetwork, resources: &ResourceRegistry) -> Result<(), PreprocessorError> {
 		let mut hash_to_node_id: HashMap<graph_craft::application_io::resource::ResourceHash, NodeId> = HashMap::new();
 		let mut new_resource_nodes: Vec<(NodeId, DocumentNode)> = Vec::new();
 
 		for node in network.nodes.values_mut() {
 			if let DocumentNodeImplementation::Network(nested) = &mut node.implementation {
-			replace_resource_inputs(nested, resources)?;
+				self.replace_resource_inputs(nested, resources)?;
 				continue;
 			}
 
@@ -72,27 +85,10 @@ fn replace_resource_inputs(network: &mut NodeNetwork, resources: &ResourceRegist
 		Ok(())
 	}
 
-impl Preprocessor {
-	fn insert_inject_scopes(&self, network: &mut NodeNetwork) {
-		for (identifier, (template, ty)) in self.inject_scopes.iter() {
-			let mut hasher = DefaultHasher::new();
-
-			identifier.as_str().hash(&mut hasher);
-			let producer_id = NodeId(hasher.finish());
-			network.nodes.insert(producer_id, template.clone());
-
-			network.scope_injections.insert(identifier.as_str().to_string(), (producer_id, ty.clone()));
-		}
-	}
-
-	fn expand_network_inner(&self, network: &mut NodeNetwork) {
-		if network.generated {
-			return;
-		}
-
+	fn expand_network(&self, network: &mut NodeNetwork) {
 		for node in network.nodes.values_mut() {
 			match &mut node.implementation {
-				DocumentNodeImplementation::Network(node_network) => self.expand_network_inner(node_network),
+				DocumentNodeImplementation::Network(node_network) => self.expand_network(node_network),
 				DocumentNodeImplementation::ProtoNode(proto_node_identifier) => {
 					if let Some(new_node) = self.substitutions.get(proto_node_identifier) {
 						// Reconcile the document node's inputs with what the current node definition expects,