Give extent overrides a typed input surface instead of the raw node form

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<T>,
+}
+
+impl<'a, T> ValueIn<'a, T> {
+	pub fn new(read: &'a dyn Fn() -> GPoll<T>) -> Self {
+		Self { read }
+	}
+
+	pub fn get(&self) -> GPoll<T> {
+		(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<Extent>,
+}
+
+impl<'a> ExtentIn<'a> {
+	pub fn new(query: &'a dyn Fn(u64, u8) -> GPoll<Extent>) -> Self {
+		Self { query }
+	}
+
+	pub fn at(&self, level: LevelIn) -> GPoll<Extent> {
+		(self.query)(0, level.level)
+	}
+
+	pub fn at_copy(&self, copy: u64, level: LevelIn) -> GPoll<Extent> {
+		(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
+	}
+}

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;

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<TokenStream2> = Vec::new();
+		let mut arg_names: Vec<Ident> = 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, <Self as #core_types::node::Node<#ctx_ident>>::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 {

node-graph/node-macro/src/parsing.rs

@@ -129,6 +129,8 @@ pub(crate) struct NodeFnAttributes {
 	pub(crate) placeholder: Option<Path>,
 	/// Function overriding the generated `extent` method
 	pub(crate) extent: Option<Path>,
+	/// Function overriding the generated `extent` method with the raw node/ctx/level form
+	pub(crate) extent_raw: Option<Path>,
 	/// Function overriding the generated `eval_batch` method
 	pub(crate) batch: Option<Path>,
 	/// 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,

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<C, NodeContent>(node: &MemoizeNode<NodeContent>, ctx: &C, level: u8) -> GPoll<Extent>
-where
-	NodeContent: Node<C>,
-{
-	node.content.extent_at(ctx, level)
+fn memoize_extent(content: ExtentIn<'_>, level: LevelIn) -> GPoll<Extent> {
+	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<C, NodeContent>(node: &FrameMemoNode<NodeContent>, ctx: &C, level: u8) -> GPoll<Extent>
-where
-	NodeContent: Node<C>,
-{
-	node.content.extent_at(ctx, level)
+fn frame_memo_extent(content: ExtentIn<'_>, level: LevelIn) -> GPoll<Extent> {
+	content.at(level)
 }
 
 type MonitorValue = Arc<Mutex<Option<IORecord<CtxSnapshot, RecordCapture>>>>;

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>) -> f64 {
 }
 
 /// The pushed level's extent is the copy count; other levels forward to the carrier.
-fn repeat_opacity_extent<C, In0, In1>(node: &RepeatOpacityNode<In0, In1>, ctx: &C, level: u8) -> core_types::gpoll::GPoll<core_types::gpoll::Extent>
-where
-	In0: core_types::node::Node<C>,
-	In1: core_types::node::Node<C, Output = u32>,
-{
-	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<Extent> {
+	match level.pushed() {
+		true => count.get().map(|count| Extent::Exactly(count as usize)),
+		false => element.at(level),
 	}
 }
 
@@ -87,21 +82,10 @@ fn repeat<T>(ctx: impl Ctx + DeriveCtx + ExtractIndex, content: impl Node<Contex
 
 /// The pushed level's extent is the copy count; inner levels forward to the
 /// content, whose extent is taken uniform across copies (queried at copy 0).
-fn repeat_extent<'r, C, In0, In1>(node: &RepeatNode<In0, In1>, ctx: &C, level: u8) -> core_types::gpoll::GPoll<core_types::gpoll::Extent>
-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<C, Output = core_types::record::RecordValue<'r>>,
-{
-	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<Extent> {
+	match level.pushed() {
+		true => count.get().map(|count| Extent::Exactly(count as usize)),
+		false => content.at(level),
 	}
 }