Add node macro support for injecting nodes into global scope (#4221)

* Macro support

* Preprocessor support

* Allow Cow str as scope name

* Use Passthrough node to work around compiler auto wrapping the producer node and invalidating the node id

* Stable NodeIds

* Less hacky but with bad performace

* Remove Cow from RegistryValueSource::Scope

* Connect scope injections to the inner network export

---------

Co-authored-by: Dennis Kobert <dennis@kobert.dev>
This commit is contained in:
Timon
2026-06-11 18:14:58 +00:00
committed by GitHub
parent 05d3d66587
commit e524dff1a7
13 changed files with 272 additions and 165 deletions

View File

@@ -1,6 +1,7 @@
#[macro_use]
extern crate log;
use graph_craft::Type;
use graph_craft::application_io::resource::{ResourceId, ResourceRegistry};
use graph_craft::document::value::*;
use graph_craft::document::*;
@@ -9,11 +10,21 @@ use graph_craft::{ProtoNodeIdentifier, concrete};
use graphene_std::registry::*;
use graphene_std::*;
use std::collections::{HashMap, HashSet};
use std::hash::{DefaultHasher, Hash, Hasher};
pub fn expand_network(network: &mut NodeNetwork, substitutions: &HashMap<ProtoNodeIdentifier, DocumentNode>, resources: &ResourceRegistry) -> Result<(), PreprocessorError> {
replace_resource_inputs(network, resources)?;
expand_network_inner(network, substitutions);
Ok(())
#[derive(Debug, Default, Clone)]
pub struct Preprocessor {
substitutions: HashMap<ProtoNodeIdentifier, DocumentNode>,
inject_scopes: HashMap<ProtoNodeIdentifier, (DocumentNode, Type)>,
}
impl Preprocessor {
pub fn expand_network(&self, network: &mut NodeNetwork, resources: &ResourceRegistry) -> Result<(), PreprocessorError> {
self.insert_inject_scopes(network);
replace_resource_inputs(network, resources)?;
self.expand_network_inner(network);
Ok(())
}
}
/// Replace every `TaggedValue::Resource(hash)` input with a reference to a freshly inserted `resource` proto node.
@@ -61,160 +72,191 @@ fn replace_resource_inputs(network: &mut NodeNetwork, resources: &ResourceRegist
Ok(())
}
fn expand_network_inner(network: &mut NodeNetwork, substitutions: &HashMap<ProtoNodeIdentifier, DocumentNode>) {
if network.generated {
return;
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()));
}
}
for node in network.nodes.values_mut() {
match &mut node.implementation {
DocumentNodeImplementation::Network(node_network) => expand_network_inner(node_network, substitutions),
DocumentNodeImplementation::ProtoNode(proto_node_identifier) => {
if let Some(new_node) = substitutions.get(proto_node_identifier) {
// Reconcile the document node's inputs with what the current node definition expects,
// since the saved document may have fewer or more inputs than the current version
while node.inputs.len() < new_node.inputs.len() {
node.inputs.push(new_node.inputs[node.inputs.len()].clone());
}
node.inputs.truncate(new_node.inputs.len());
fn expand_network_inner(&self, network: &mut NodeNetwork) {
if network.generated {
return;
}
node.implementation = new_node.implementation.clone();
for node in network.nodes.values_mut() {
match &mut node.implementation {
DocumentNodeImplementation::Network(node_network) => self.expand_network_inner(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,
// since the saved document may have fewer or more inputs than the current version
while node.inputs.len() < new_node.inputs.len() {
node.inputs.push(new_node.inputs[node.inputs.len()].clone());
}
node.inputs.truncate(new_node.inputs.len());
node.implementation = new_node.implementation.clone();
}
}
DocumentNodeImplementation::Extract => (),
}
}
}
pub fn new() -> Self {
let mut substitutions = HashMap::new();
let mut inject_scopes = HashMap::new();
// We pre initialize the node registry here to avoid a deadlock
let into_node_registry = &*interpreted_executor::node_registry::NODE_REGISTRY;
let node_registry = core_types::registry::NODE_REGISTRY.lock().unwrap();
for (id, metadata) in core_types::registry::NODE_METADATA.lock().unwrap().iter() {
let id = id.clone();
let NodeMetadata { fields, memoize, inject_scope, .. } = metadata;
let Some(implementations) = node_registry.get(&id) else { continue };
let valid_call_args: HashSet<_> = implementations.iter().map(|(_, node_io)| node_io.call_argument.clone()).collect();
let first_node_io = implementations.first().map(|(_, node_io)| node_io).unwrap_or(const { &NodeIOTypes::empty() });
let mut node_io_types = vec![HashSet::new(); fields.len()];
for (_, node_io) in implementations.iter() {
for (i, ty) in node_io.inputs.iter().enumerate() {
node_io_types[i].insert(ty.clone());
}
}
DocumentNodeImplementation::Extract => (),
}
}
}
pub fn generate_node_substitutions() -> HashMap<ProtoNodeIdentifier, DocumentNode> {
let mut custom = HashMap::new();
// We pre initialize the node registry here to avoid a deadlock
let into_node_registry = &*interpreted_executor::node_registry::NODE_REGISTRY;
let node_registry = core_types::registry::NODE_REGISTRY.lock().unwrap();
for (id, metadata) in core_types::registry::NODE_METADATA.lock().unwrap().iter() {
let id = id.clone();
let NodeMetadata { fields, memoize, .. } = metadata;
let Some(implementations) = node_registry.get(&id) else { continue };
let valid_call_args: HashSet<_> = implementations.iter().map(|(_, node_io)| node_io.call_argument.clone()).collect();
let first_node_io = implementations.first().map(|(_, node_io)| node_io).unwrap_or(const { &NodeIOTypes::empty() });
let mut node_io_types = vec![HashSet::new(); fields.len()];
for (_, node_io) in implementations.iter() {
for (i, ty) in node_io.inputs.iter().enumerate() {
node_io_types[i].insert(ty.clone());
let mut input_type = &first_node_io.call_argument;
if valid_call_args.len() > 1 {
input_type = &const { generic!(D) };
}
}
let mut input_type = &first_node_io.call_argument;
if valid_call_args.len() > 1 {
input_type = &const { generic!(D) };
}
let inputs: Vec<_> = node_inputs(fields, first_node_io);
let input_count = inputs.len();
let network_inputs = (0..input_count).map(|i| NodeInput::node(NodeId(i as u64), 0)).collect();
let inputs: Vec<_> = node_inputs(fields, first_node_io);
let input_count = inputs.len();
let network_inputs = (0..input_count).map(|i| NodeInput::node(NodeId(i as u64), 0)).collect();
let passthrough_node = ops::passthrough::IDENTIFIER;
let passthrough_node = ops::passthrough::IDENTIFIER;
let mut generated_nodes = 0;
let mut nodes: HashMap<_, _, _> = node_io_types
.iter()
.take(input_count)
.enumerate()
.map(|(i, inputs)| {
(
NodeId(i as u64),
match inputs.len() {
1 => {
let input = inputs.iter().next().unwrap();
let input_ty = input.nested_type();
let mut inputs = vec![NodeInput::import(input.clone(), i)];
let mut generated_nodes = 0;
let mut nodes: HashMap<_, _, _> = node_io_types
.iter()
.take(input_count)
.enumerate()
.map(|(i, inputs)| {
(
NodeId(i as u64),
match inputs.len() {
1 => {
let input = inputs.iter().next().unwrap();
let input_ty = input.nested_type();
let mut inputs = vec![NodeInput::import(input.clone(), i)];
let into_node_identifier = ProtoNodeIdentifier::with_owned_string(format!("graphene_core::ops::IntoNode<{}>", input_ty.identifier_name()));
let convert_node_identifier = ProtoNodeIdentifier::with_owned_string(format!("graphene_core::ops::ConvertNode<{}>", input_ty.identifier_name()));
let into_node_identifier = ProtoNodeIdentifier::with_owned_string(format!("graphene_core::ops::IntoNode<{}>", input_ty.identifier_name()));
let convert_node_identifier = ProtoNodeIdentifier::with_owned_string(format!("graphene_core::ops::ConvertNode<{}>", input_ty.identifier_name()));
let proto_node = if into_node_registry.keys().any(|ident: &ProtoNodeIdentifier| ident.as_str() == into_node_identifier.as_str()) {
generated_nodes += 1;
into_node_identifier
} else if into_node_registry.keys().any(|ident| ident.as_str() == convert_node_identifier.as_str()) {
generated_nodes += 1;
inputs.push(NodeInput::value(TaggedValue::None, false));
convert_node_identifier
} else {
passthrough_node.clone()
};
let mut original_location = OriginalLocation::default();
original_location.auto_convert_index = Some(i);
DocumentNode {
inputs,
implementation: DocumentNodeImplementation::ProtoNode(proto_node),
visible: true,
original_location,
..Default::default()
let proto_node = if into_node_registry.keys().any(|ident: &ProtoNodeIdentifier| ident.as_str() == into_node_identifier.as_str()) {
generated_nodes += 1;
into_node_identifier
} else if into_node_registry.keys().any(|ident| ident.as_str() == convert_node_identifier.as_str()) {
generated_nodes += 1;
inputs.push(NodeInput::value(TaggedValue::None, false));
convert_node_identifier
} else {
passthrough_node.clone()
};
let mut original_location = OriginalLocation::default();
original_location.auto_convert_index = Some(i);
DocumentNode {
inputs,
implementation: DocumentNodeImplementation::ProtoNode(proto_node),
visible: true,
original_location,
..Default::default()
}
}
}
_ => DocumentNode {
inputs: vec![NodeInput::import(generic!(X), i)],
implementation: DocumentNodeImplementation::ProtoNode(passthrough_node.clone()),
visible: false,
..Default::default()
_ => DocumentNode {
inputs: vec![NodeInput::import(generic!(X), i)],
implementation: DocumentNodeImplementation::ProtoNode(passthrough_node.clone()),
visible: false,
..Default::default()
},
},
},
)
})
.collect();
)
})
.collect();
if generated_nodes == 0 && !memoize {
continue;
}
if generated_nodes == 0 && !memoize && !inject_scope {
continue;
}
let document_node = DocumentNode {
inputs: network_inputs,
call_argument: input_type.clone(),
implementation: DocumentNodeImplementation::ProtoNode(id.clone()),
visible: true,
skip_deduplication: false,
context_features: ContextDependencies::from(metadata.context_features.as_slice()),
..Default::default()
};
nodes.insert(NodeId(input_count as u64), document_node);
// If memoize is requested, append a Memoize node after the main node and redirect the export through it
let export_node_id = if *memoize {
let memoize_node_id = NodeId(input_count as u64 + 1);
let memoize_node = DocumentNode {
inputs: vec![NodeInput::node(NodeId(input_count as u64), 0)],
implementation: DocumentNodeImplementation::ProtoNode(graphene_core::memo::memoize::IDENTIFIER.clone()),
let document_node = DocumentNode {
inputs: network_inputs,
call_argument: input_type.clone(),
implementation: DocumentNodeImplementation::ProtoNode(id.clone()),
visible: true,
skip_deduplication: false,
context_features: ContextDependencies::from(metadata.context_features.as_slice()),
..Default::default()
};
nodes.insert(memoize_node_id, memoize_node);
memoize_node_id
} else {
NodeId(input_count as u64)
};
let node = DocumentNode {
inputs,
call_argument: input_type.clone(),
implementation: DocumentNodeImplementation::Network(NodeNetwork {
exports: vec![NodeInput::Node {
node_id: export_node_id,
output_index: 0,
}],
nodes,
scope_injections: Default::default(),
generated: true,
}),
visible: true,
skip_deduplication: false,
..Default::default()
};
nodes.insert(NodeId(input_count as u64), document_node);
custom.insert(id.clone(), node);
// If memoize is requested, append a Memoize node after the main node and redirect the export through it
let export_node_id = if *memoize {
let memoize_node_id = NodeId(input_count as u64 + 1);
let memoize_node = DocumentNode {
inputs: vec![NodeInput::node(NodeId(input_count as u64), 0)],
implementation: DocumentNodeImplementation::ProtoNode(graphene_core::memo::memoize::IDENTIFIER.clone()),
visible: true,
..Default::default()
};
nodes.insert(memoize_node_id, memoize_node);
memoize_node_id
} else {
NodeId(input_count as u64)
};
let node = DocumentNode {
inputs,
call_argument: input_type.clone(),
implementation: DocumentNodeImplementation::Network(NodeNetwork {
exports: vec![NodeInput::Node {
node_id: export_node_id,
output_index: 0,
}],
nodes,
scope_injections: Default::default(),
generated: true,
}),
visible: true,
skip_deduplication: false,
..Default::default()
};
substitutions.insert(id.clone(), node);
// If `inject_scope` is requested, prepare the proto node template and type info needed
if *inject_scope
&& let Some(implementations) = node_registry.get(&id)
&& let Some((_, node_io)) = implementations.first()
{
let template = DocumentNode {
inputs: node_inputs(fields, node_io),
call_argument: node_io.call_argument.clone(),
implementation: DocumentNodeImplementation::ProtoNode(id.clone()),
visible: true,
context_features: ContextDependencies::from(metadata.context_features.as_slice()),
..Default::default()
};
inject_scopes.insert(id.clone(), (template, node_io.return_value.clone()));
}
}
Self { substitutions, inject_scopes }
}
custom
}
pub fn node_inputs(fields: &[registry::FieldMetadata], first_node_io: &NodeIOTypes) -> Vec<NodeInput> {
@@ -229,7 +271,7 @@ pub fn node_inputs(fields: &[registry::FieldMetadata], first_node_io: &NodeIOTyp
};
let exposed = if index == 0 { *ty != fn_type_fut!(Context, ()) } else { field.exposed };
match field.value_source {
match &field.value_source {
RegistryValueSource::None => {}
RegistryValueSource::Default(data) => {
if let Some(custom_default) = TaggedValue::from_primitive_string(data, ty) {
@@ -239,7 +281,7 @@ pub fn node_inputs(fields: &[registry::FieldMetadata], first_node_io: &NodeIOTyp
warn!("Failed to parse default value for type `{ty:?}` with data `{data}`");
}
}
RegistryValueSource::Scope(data) => return NodeInput::scope(Cow::Borrowed(data)),
RegistryValueSource::Scope(data) => return NodeInput::scope(*data),
};
if let Some(type_default) = TaggedValue::from_type(ty) {