From 131abf5883c29b93d1113972430d5f07b85514a7 Mon Sep 17 00:00:00 2001 From: Dennis Kobert Date: Sun, 16 Aug 2026 15:55:17 +0000 Subject: [PATCH] Give extent overrides a typed input surface instead of the raw node form --- node-graph/libraries/core-types/src/extent.rs | 61 ++++++++++++++++ node-graph/libraries/core-types/src/lib.rs | 1 + node-graph/node-macro/src/codegen.rs | 69 +++++++++++++++++-- node-graph/node-macro/src/parsing.rs | 30 +++++++- node-graph/nodes/gcore/src/memo.rs | 15 ++-- node-graph/nodes/gcore/src/record.rs | 36 +++------- 6 files changed, 171 insertions(+), 41 deletions(-) create mode 100644 node-graph/libraries/core-types/src/extent.rs diff --git a/node-graph/libraries/core-types/src/extent.rs b/node-graph/libraries/core-types/src/extent.rs new file mode 100644 index 0000000000..e00830e629 --- /dev/null +++ b/node-graph/libraries/core-types/src/extent.rs @@ -0,0 +1,61 @@ +//! The typed surface handed to an `extent(fn)` helper: the node's inputs in +//! declaration order, then the queried level. Values read without unsafe or +//! internal fields, upstream extents query per level, and the one blessed +//! context modification is per-copy derived promotion. Anything beyond this +//! vocabulary uses `extent_raw(fn)`, which keeps the full node/ctx/level form. + +use crate::gpoll::{Extent, GPoll}; + +/// A wired value input; `get` evaluates the edge and yields the typed element. +pub struct ValueIn<'a, T> { + read: &'a dyn Fn() -> GPoll, +} + +impl<'a, T> ValueIn<'a, T> { + pub fn new(read: &'a dyn Fn() -> GPoll) -> Self { + Self { read } + } + + pub fn get(&self) -> GPoll { + (self.read)() + } +} + +/// An upstream edge's extents. For derived (per-copy) content the query runs +/// at the given copy's promoted context; `at` queries copy 0, the uniform +/// default. For ordinary edges the copy is ignored. +pub struct ExtentIn<'a> { + query: &'a dyn Fn(u64, u8) -> GPoll, +} + +impl<'a> ExtentIn<'a> { + pub fn new(query: &'a dyn Fn(u64, u8) -> GPoll) -> Self { + Self { query } + } + + pub fn at(&self, level: LevelIn) -> GPoll { + (self.query)(0, level.level) + } + + pub fn at_copy(&self, copy: u64, level: LevelIn) -> GPoll { + (self.query)(copy, level.level) + } +} + +/// The queried absolute level (innermost `0`), paired with the node's depth. +#[derive(Clone, Copy, Debug)] +pub struct LevelIn { + pub level: u8, + pub depth: u8, +} + +impl LevelIn { + pub fn new(level: u8, depth: u8) -> Self { + Self { level, depth } + } + + /// Whether the query targets the node's own pushed (outermost) level. + pub fn pushed(&self) -> bool { + self.level + 1 == self.depth + } +} diff --git a/node-graph/libraries/core-types/src/lib.rs b/node-graph/libraries/core-types/src/lib.rs index ff702baa5a..20a95eeab5 100644 --- a/node-graph/libraries/core-types/src/lib.rs +++ b/node-graph/libraries/core-types/src/lib.rs @@ -5,6 +5,7 @@ pub mod attribute; pub mod bounds; pub mod consts; pub mod context; +pub mod extent; pub mod frame_table; pub mod gpoll; pub mod list; diff --git a/node-graph/node-macro/src/codegen.rs b/node-graph/node-macro/src/codegen.rs index f3a28610c1..beeb547700 100644 --- a/node-graph/node-macro/src/codegen.rs +++ b/node-graph/node-macro/src/codegen.rs @@ -1221,13 +1221,74 @@ pub(crate) fn generate_node_impl(crate_ident: &CrateIdent, parsed: &ParsedNodeFn // The extent override is the leveled `extent_at`; consumers query the // composite `extent(ctx, Level)`, which the trait derives from it. A node // without `extent = fn` keeps the scalar default (one item at every level). - let extent_impl = match &parsed.attributes.extent { - Some(path) => quote! { + // The typed extent surface: the node's inputs in declaration order (values + // readable without unsafe, edges as per-level extent queries, derived + // content promoted per copy), then the level paired with the node's depth. + let extent_impl = if let Some(path) = &parsed.attributes.extent { + let mut arg_decls: Vec = Vec::new(); + let mut arg_names: Vec = Vec::new(); + for (index, field) in regular_fields.iter().enumerate() { + let name = &field.pat_ident.ident; + if injected_name(name) { + continue; + } + let arg = format_ident!("__extent_arg_{index}"); + let query = format_ident!("__extent_query_{index}"); + let extent_edge = |query: &Ident, arg: &Ident| { + quote! { + let #query = |_: u64, __lvl: u8| #core_types::node::Node::extent_at(&self.#name, __input, __lvl); + let #arg = #core_types::extent::ExtentIn::new(&#query); + } + }; + let decl = match &field.ty { + ParsedFieldType::Node(_) => match ir::lazy_binding(&node, index) { + ir::LazyBinding::DeriveRouting => quote! { + let #query = |__copy: u64, __lvl: u8| { + let __head = #core_types::context::DeriveCtx::index_head(__input); + #core_types::record::DerivedRecordEdge::extent_at_derived(&self.#name, &#core_types::context::DeriveCtx::promoted(__input, &__head, __copy), __lvl) + }; + let #arg = #core_types::extent::ExtentIn::new(&#query); + }, + _ => extent_edge(&query, &arg), + }, + ParsedFieldType::Regular(RegularParsedField { ty, .. }) => match ir::value_binding(&node, index) { + ValueBinding::RecordElement | ValueBinding::ReadingSecondary => { + let slot = format_ident!("__in_{index}"); + quote! { + let #query = || { + #core_types::node::Node::eval(&self.#name, __input) + .map(|__value| unsafe { #core_types::record::read_element::<#ty>(self.#slot.rec(&__value)) }) + }; + let #arg = #core_types::extent::ValueIn::new(&#query); + } + } + ValueBinding::Plain => quote! { + let #query = || #core_types::node::Node::eval(&self.#name, __input); + let #arg = #core_types::extent::ValueIn::new(&#query); + }, + // A carrier, lent, or materialized ranked input is a record + // edge; its extents are the queryable quantity. + _ => extent_edge(&query, &arg), + }, + }; + arg_decls.push(decl); + arg_names.push(arg); + } + quote! { + fn extent_at(&self, __input: &#ctx_ident, __level: u8) -> #core_types::gpoll::GPoll<#core_types::gpoll::Extent> { + #(#arg_decls)* + let __level_in = #core_types::extent::LevelIn::new(__level, >::layout(self).depth); + #path(#(#arg_names,)* __level_in) + } + } + } else if let Some(path) = &parsed.attributes.extent_raw { + quote! { fn extent_at(&self, __input: &#ctx_ident, __level: u8) -> #core_types::gpoll::GPoll<#core_types::gpoll::Extent> { #path(self, __input, __level) } - }, - None => quote!(), + } + } else { + quote!() }; let serialize_impl = match &parsed.attributes.serialize { diff --git a/node-graph/node-macro/src/parsing.rs b/node-graph/node-macro/src/parsing.rs index b751bbe265..70708e0d60 100644 --- a/node-graph/node-macro/src/parsing.rs +++ b/node-graph/node-macro/src/parsing.rs @@ -129,6 +129,8 @@ pub(crate) struct NodeFnAttributes { pub(crate) placeholder: Option, /// Function overriding the generated `extent` method pub(crate) extent: Option, + /// Function overriding the generated `extent` method with the raw node/ctx/level form + pub(crate) extent_raw: Option, /// Function overriding the generated `eval_batch` method pub(crate) batch: Option, /// Whether partial upstream values are mapped to `Pending` instead of flowing into this node @@ -400,6 +402,7 @@ impl Parse for NodeFnAttributes { let mut inject_scope = false; let mut placeholder = None; let mut extent = None; + let mut extent_raw = None; let mut batch = None; let mut no_partial = false; let mut plain = false; @@ -572,6 +575,19 @@ impl Parse for NodeFnAttributes { let parsed_path: Path = meta.parse_args().map_err(|_| Error::new_spanned(meta, "Expected a valid path for 'extent', e.g., extent(my_extent)"))?; extent = Some(parsed_path); } + // Escape hatch for extent overrides needing arbitrary context access: the raw + // `(node, ctx, level)` form instead of the typed `extent(fn)` input surface. + // + // Example usage: + // #[node_macro::node(..., extent_raw(my_extent), ...)] + "extent_raw" => { + let meta = meta.require_list()?; + if extent_raw.is_some() { + return Err(Error::new_spanned(meta, "Multiple 'extent_raw' attributes are not allowed")); + } + let parsed_path: Path = meta.parse_args().map_err(|_| Error::new_spanned(meta, "Expected a valid path for 'extent_raw', e.g., extent_raw(my_extent)"))?; + extent_raw = Some(parsed_path); + } // Function overriding the generated `eval_batch` method, replacing the trait's per-lane spec loop. // // Example usage: @@ -608,7 +624,7 @@ impl Parse for NodeFnAttributes { indoc!( r#" Unsupported attribute in `node`. - Supported attributes are 'category', 'name', 'path', 'skip_impl', 'properties', 'cfg', 'shader_node', 'serialize', 'memoize', 'inject_scope', 'placeholder', 'extent', 'batch', and 'no_partial'. + Supported attributes are 'category', 'name', 'path', 'skip_impl', 'properties', 'cfg', 'shader_node', 'serialize', 'memoize', 'inject_scope', 'placeholder', 'extent', 'extent_raw', 'batch', and 'no_partial'. Example usage: #[node_macro::node(..., name("Test Node"), ...)] "# @@ -631,6 +647,10 @@ impl Parse for NodeFnAttributes { )); } + if let (Some(_), Some(raw)) = (&extent, &extent_raw) { + return Err(Error::new_spanned(raw, "'extent' and 'extent_raw' are mutually exclusive")); + } + Ok(NodeFnAttributes { category, display_name, @@ -644,6 +664,7 @@ impl Parse for NodeFnAttributes { inject_scope, placeholder, extent, + extent_raw, batch, no_partial, plain, @@ -1393,6 +1414,7 @@ mod tests { inject_scope: false, placeholder: None, extent: None, + extent_raw: None, batch: None, no_partial: false, plain: false, @@ -1472,6 +1494,7 @@ mod tests { inject_scope: false, placeholder: None, extent: None, + extent_raw: None, batch: None, no_partial: false, plain: false, @@ -1566,6 +1589,7 @@ mod tests { inject_scope: false, placeholder: None, extent: None, + extent_raw: None, batch: None, no_partial: false, plain: false, @@ -1641,6 +1665,7 @@ mod tests { inject_scope: false, placeholder: None, extent: None, + extent_raw: None, batch: None, no_partial: false, plain: false, @@ -1728,6 +1753,7 @@ mod tests { inject_scope: false, placeholder: None, extent: None, + extent_raw: None, batch: None, no_partial: false, plain: false, @@ -1818,6 +1844,7 @@ mod tests { inject_scope: false, placeholder: None, extent: None, + extent_raw: None, batch: None, no_partial: false, plain: false, @@ -1893,6 +1920,7 @@ mod tests { inject_scope: false, placeholder: None, extent: None, + extent_raw: None, batch: None, no_partial: false, plain: false, diff --git a/node-graph/nodes/gcore/src/memo.rs b/node-graph/nodes/gcore/src/memo.rs index a7ac02527a..ed3d698609 100644 --- a/node-graph/nodes/gcore/src/memo.rs +++ b/node-graph/nodes/gcore/src/memo.rs @@ -1,5 +1,6 @@ use core_types::arena::ArenaCell; use core_types::context::{Ctx, CtxSnapshot, DeriveCtx, ExtractAll, ExtractArena}; +use core_types::extent::{ExtentIn, LevelIn}; use core_types::frame_table::{FrameTable, Lookup}; use core_types::gpoll::{Extent, Finality, GPoll}; use core_types::graphene_hash::CacheHash; @@ -45,11 +46,8 @@ fn memoize<'e>( result } -fn memoize_extent(node: &MemoizeNode, ctx: &C, level: u8) -> GPoll -where - NodeContent: Node, -{ - node.content.extent_at(ctx, level) +fn memoize_extent(content: ExtentIn<'_>, level: LevelIn) -> GPoll { + content.at(level) } #[node_macro::node(category(""), path(graphene_core::memo), extent(frame_memo_extent))] @@ -95,11 +93,8 @@ fn frame_memo<'e>( } } -fn frame_memo_extent(node: &FrameMemoNode, ctx: &C, level: u8) -> GPoll -where - NodeContent: Node, -{ - node.content.extent_at(ctx, level) +fn frame_memo_extent(content: ExtentIn<'_>, level: LevelIn) -> GPoll { + content.at(level) } type MonitorValue = Arc>>>; diff --git a/node-graph/nodes/gcore/src/record.rs b/node-graph/nodes/gcore/src/record.rs index e527b32e89..ff3b10a6b5 100644 --- a/node-graph/nodes/gcore/src/record.rs +++ b/node-graph/nodes/gcore/src/record.rs @@ -7,7 +7,8 @@ use core_types::attribute::{Attr, Opacity, RemoveAttr}; use core_types::context::{DeriveCtx, ExtractArena, ExtractIndex, InjectIndex}; -use core_types::gpoll::{ErrorKind, GraphError, Interrupt}; +use core_types::extent::{ExtentIn, LevelIn, ValueIn}; +use core_types::gpoll::{ErrorKind, Extent, GPoll, GraphError, Interrupt}; use core_types::{Context, Ctx}; core_types::attribute! { @@ -63,16 +64,10 @@ fn sum(_: impl Ctx + InjectIndex + Copy, items: IList) -> f64 { } /// The pushed level's extent is the copy count; other levels forward to the carrier. -fn repeat_opacity_extent(node: &RepeatOpacityNode, ctx: &C, level: u8) -> core_types::gpoll::GPoll -where - In0: core_types::node::Node, - In1: core_types::node::Node, -{ - use core_types::node::Node; - if level + 1 == node.__layout.depth { - node.count.eval(ctx).map(|count| core_types::gpoll::Extent::Exactly(count as usize)) - } else { - node.element.extent_at(ctx, level) +fn repeat_opacity_extent(element: ExtentIn<'_>, count: ValueIn<'_, u32>, level: LevelIn) -> GPoll { + match level.pushed() { + true => count.get().map(|count| Extent::Exactly(count as usize)), + false => element.at(level), } } @@ -87,21 +82,10 @@ fn repeat(ctx: impl Ctx + DeriveCtx + ExtractIndex, content: impl Node(node: &RepeatNode, ctx: &C, level: u8) -> core_types::gpoll::GPoll -where - C: core_types::context::DeriveCtx + core_types::context::ExtractIndex, - In0: for<'d> core_types::record::DerivedRecordEdge<'d, core_types::context::Derived<'d, C>>, - In1: core_types::node::Node>, -{ - use core_types::node::Node; - if level + 1 == node.__layout.depth { - node.count.eval(ctx).map(|value| { - let count: u32 = unsafe { core_types::record::read_element(node.__in_1.rec(&value)) }; - core_types::gpoll::Extent::Exactly(count as usize) - }) - } else { - let spilled = ctx.index_head(); - node.content.extent_at_derived(&ctx.promoted(&spilled, 0), level) +fn repeat_extent(content: ExtentIn<'_>, count: ValueIn<'_, u32>, level: LevelIn) -> GPoll { + match level.pushed() { + true => count.get().map(|count| Extent::Exactly(count as usize)), + false => content.at(level), } }