From e0368435b95a9edcfdc6ea48d09596262faed9db Mon Sep 17 00:00:00 2001 From: Dennis Kobert Date: Tue, 28 Apr 2026 13:55:38 +0200 Subject: [PATCH] Add the "memoize" attribute to the node macro (#4065) * Add memoization attribute to node macro * Fix memoization insertion for networks without conversion nodes --- .../document_node_derive.rs | 1 + .../libraries/core-types/src/registry.rs | 1 + node-graph/node-macro/src/codegen.rs | 2 ++ node-graph/node-macro/src/parsing.rs | 26 +++++++++++++++++-- node-graph/nodes/vector/src/vector_nodes.rs | 2 +- node-graph/preprocessor/src/lib.rs | 21 ++++++++++++--- 6 files changed, 47 insertions(+), 6 deletions(-) diff --git a/editor/src/messages/portfolio/document/node_graph/document_node_definitions/document_node_derive.rs b/editor/src/messages/portfolio/document/node_graph/document_node_definitions/document_node_derive.rs index 72a4558a..4d0c14fd 100644 --- a/editor/src/messages/portfolio/document/node_graph/document_node_definitions/document_node_derive.rs +++ b/editor/src/messages/portfolio/document/node_graph/document_node_definitions/document_node_derive.rs @@ -36,6 +36,7 @@ pub(super) fn post_process_nodes(custom: Vec) -> HashMap description, properties, context_features, + memoize: _, } = metadata; let Some(implementations) = &node_registry.get(id) else { continue }; diff --git a/node-graph/libraries/core-types/src/registry.rs b/node-graph/libraries/core-types/src/registry.rs index a472ef45..9e9f1267 100644 --- a/node-graph/libraries/core-types/src/registry.rs +++ b/node-graph/libraries/core-types/src/registry.rs @@ -16,6 +16,7 @@ pub struct NodeMetadata { pub description: &'static str, pub properties: Option<&'static str>, pub context_features: Vec, + pub memoize: bool, } // Translation struct between macro and definition diff --git a/node-graph/node-macro/src/codegen.rs b/node-graph/node-macro/src/codegen.rs index 49d49b7a..da2c0537 100644 --- a/node-graph/node-macro/src/codegen.rs +++ b/node-graph/node-macro/src/codegen.rs @@ -401,6 +401,7 @@ pub(crate) fn generate_node_code(crate_ident: &CrateIdent, parsed: &ParsedNodeFn let import_name = format_ident!("_IMPORT_STUB_{}", mod_name.to_string().to_case(Case::UpperSnake)); let properties = &attributes.properties_string.as_ref().map(|value| quote!(Some(#value))).unwrap_or(quote!(None)); + let memoize_flag = attributes.memoize; let cfg = crate::shader_nodes::modify_cfg(attributes); let node_input_accessor = generate_node_input_references(parsed, fn_generics, &field_idents, core_types, &identifier, &cfg); @@ -498,6 +499,7 @@ pub(crate) fn generate_node_code(crate_ident: &CrateIdent, parsed: &ParsedNodeFn description: #description, properties: #properties, context_features: vec![#(ContextFeature::#context_features,)*], + memoize: #memoize_flag, fields: vec![ #( FieldMetadata { diff --git a/node-graph/node-macro/src/parsing.rs b/node-graph/node-macro/src/parsing.rs index ac140b85..329fb61c 100644 --- a/node-graph/node-macro/src/parsing.rs +++ b/node-graph/node-macro/src/parsing.rs @@ -52,7 +52,8 @@ pub(crate) struct NodeFnAttributes { pub(crate) shader_node: Option, /// Custom serialization function path (e.g., "my_module::custom_serialize") pub(crate) serialize: Option, - // Add more attributes as needed + /// Whether the preprocessor should add a Memo node after this node in the generated subnetwork + pub(crate) memoize: bool, } #[derive(Clone, Debug, Default)] @@ -259,6 +260,7 @@ impl Parse for NodeFnAttributes { let mut cfg = None; let mut shader_node = None; let mut serialize = None; + let mut memoize = false; let content = input; // let content; @@ -377,13 +379,25 @@ impl Parse for NodeFnAttributes { .map_err(|_| Error::new_spanned(meta, "Expected a valid path for 'serialize', e.g., serialize(my_module::custom_serialize)"))?; serialize = Some(parsed_path); } + // Instructs the preprocessor to insert a Memo node after this node in the generated subnetwork, + // caching its output across evaluations with identical inputs. + // + // Example usage: + // #[node_macro::node(..., memoize, ...)] + "memoize" => { + let path = meta.require_path_only()?; + if memoize { + return Err(Error::new_spanned(path, "Multiple 'memoize' attributes are not allowed")); + } + memoize = true; + } _ => { return Err(Error::new_spanned( meta, indoc!( r#" Unsupported attribute in `node`. - Supported attributes are 'category', 'name', 'path', 'skip_impl', 'properties', 'cfg', 'shader_node', and 'serialize'. + Supported attributes are 'category', 'name', 'path', 'skip_impl', 'properties', 'cfg', 'shader_node', 'serialize', and 'memoize'. Example usage: #[node_macro::node(..., name("Test Node"), ...)] "# @@ -415,6 +429,7 @@ impl Parse for NodeFnAttributes { cfg, shader_node, serialize, + memoize, }) } } @@ -1020,6 +1035,7 @@ mod tests { cfg: None, shader_node: None, serialize: None, + memoize: false, }, fn_name: Ident::new("add", Span::call_site()), struct_name: Ident::new("Add", Span::call_site()), @@ -1088,6 +1104,7 @@ mod tests { cfg: None, shader_node: None, serialize: None, + memoize: false, }, fn_name: Ident::new("transform", Span::call_site()), struct_name: Ident::new("Transform", Span::call_site()), @@ -1170,6 +1187,7 @@ mod tests { cfg: None, shader_node: None, serialize: None, + memoize: false, }, fn_name: Ident::new("circle", Span::call_site()), struct_name: Ident::new("Circle", Span::call_site()), @@ -1234,6 +1252,7 @@ mod tests { cfg: None, shader_node: None, serialize: None, + memoize: false, }, fn_name: Ident::new("levels", Span::call_site()), struct_name: Ident::new("Levels", Span::call_site()), @@ -1310,6 +1329,7 @@ mod tests { cfg: None, shader_node: None, serialize: None, + memoize: false, }, fn_name: Ident::new("add", Span::call_site()), struct_name: Ident::new("Add", Span::call_site()), @@ -1374,6 +1394,7 @@ mod tests { cfg: None, shader_node: None, serialize: None, + memoize: false, }, fn_name: Ident::new("load_image", Span::call_site()), struct_name: Ident::new("LoadImage", Span::call_site()), @@ -1438,6 +1459,7 @@ mod tests { cfg: None, shader_node: None, serialize: None, + memoize: false, }, fn_name: Ident::new("custom_node", Span::call_site()), struct_name: Ident::new("CustomNode", Span::call_site()), diff --git a/node-graph/nodes/vector/src/vector_nodes.rs b/node-graph/nodes/vector/src/vector_nodes.rs index 82292733..4a7c1c36 100644 --- a/node-graph/nodes/vector/src/vector_nodes.rs +++ b/node-graph/nodes/vector/src/vector_nodes.rs @@ -1326,7 +1326,7 @@ pub async fn flatten_path(_: impl Ctx, #[implem } /// Convert vector geometry into a polyline composed of evenly spaced points. -#[node_macro::node(category("Vector: Modifier"), path(core_types::vector), properties("sample_polyline_properties"))] +#[node_macro::node(category("Vector: Modifier"), path(core_types::vector), properties("sample_polyline_properties"), memoize)] async fn sample_polyline( _: impl Ctx, content: Table, diff --git a/node-graph/preprocessor/src/lib.rs b/node-graph/preprocessor/src/lib.rs index f17760e6..961c084f 100644 --- a/node-graph/preprocessor/src/lib.rs +++ b/node-graph/preprocessor/src/lib.rs @@ -42,7 +42,7 @@ pub fn generate_node_substitutions() -> HashMap = 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() }); @@ -111,7 +111,7 @@ pub fn generate_node_substitutions() -> HashMap HashMap