Add the sync-prologue async source form

2026-09-15 22:28:10 +08:00 · 2026-07-27 10:47:44 +00:00
parent 470ffdab43
commit 6392c02dee
5 changed files with 244 additions and 60 deletions
--- a/node-graph/libraries/core-types/src/runtime.rs
+++ b/node-graph/libraries/core-types/src/runtime.rs
@@ -4,9 +4,9 @@ use std::pin::Pin;
 use std::sync::Arc;

 #[cfg(not(target_family = "wasm"))]
-pub type SourceFuture = Pin<Box<dyn Future<Output = ()> + Send + 'static>>;
+pub type SourceFuture<T = ()> = Pin<Box<dyn Future<Output = T> + Send + 'static>>;
 #[cfg(target_family = "wasm")]
-pub type SourceFuture = Pin<Box<dyn Future<Output = ()> + 'static>>;
+pub type SourceFuture<T = ()> = Pin<Box<dyn Future<Output = T> + 'static>>;

 #[cfg(not(target_family = "wasm"))]
 pub type DynRuntime = dyn Runtime + Send + Sync;
@@ -34,7 +34,7 @@ impl graphene_hash::CacheHash for RuntimeHandle {
 mod tests {
 	use super::*;
 	use crate::arena::Arena;
-	use crate::context::{ContextImpl, Ctx, CtxSnapshot, EvalScope, ExtractFootprint};
+	use crate::context::{ContextImpl, Ctx, CtxSnapshot, EvalScope, ExtractFootprint, ExtractVarArgs, VarArgLink, VarArgSlots};
 	use crate::gnode::GNode;
 	use crate::gpoll::GPoll;
 	use crate::transform::Footprint;
@@ -103,6 +103,38 @@ mod tests {
 		ctx.try_footprint().map(|footprint| footprint.resolution.x).unwrap_or(0)
 	}

+	#[node_macro::node(category(""))]
+	async fn snapshot_vararg(ctx: CtxSnapshot, _primary: ()) -> f64 {
+		ctx.vararg(0).ok().and_then(|slot| slot.downcast_ref::<f64>()).copied().unwrap_or(0.)
+	}
+
+	static STAGED_RUNS: AtomicU32 = AtomicU32::new(0);
+
+	#[node_macro::node(category(""))]
+	fn staged_double(_: impl Ctx, value: f64) -> SourceFuture<f64> {
+		STAGED_RUNS.fetch_add(1, Ordering::Relaxed);
+		Box::pin(async move { value * 2. })
+	}
+
+	#[node_macro::node(category(""))]
+	fn staged_sum(ctx: impl Ctx, value: f64, addend: impl Node<Context<'_>, Output = f64>) -> Result<SourceFuture<f64>, crate::gpoll::Interrupt> {
+		let addend = addend.eval(ctx)?;
+		Ok(Box::pin(async move { value + addend }))
+	}
+
+	struct GatedSource(Arc<std::sync::atomic::AtomicBool>, f64);
+
+	impl<Input> GNode<Input> for GatedSource {
+		type Output = f64;
+
+		fn eval(&self, _input: &Input) -> GPoll<f64> {
+			match self.0.load(Ordering::Relaxed) {
+				true => GPoll::Final(self.1),
+				false => GPoll::Pending,
+			}
+		}
+	}
+
 	fn scope_fixture<'a>(generations: &'a [(SourceId, u64)], arena: &'a Arena) -> EvalScope<'a> {
 		EvalScope::new(None, None, None, generations, arena)
 	}
@@ -157,6 +189,70 @@ mod tests {
 		assert_eq!(GNode::eval(&graph, &ctx), GPoll::Final(42.0));
 	}

+	#[test]
+	fn prologue_runs_sync_and_spawns_once() {
+		let arena = Arena::new(64);
+		let generations = [];
+		let scope = scope_fixture(&generations, &arena);
+		let ctx = ContextImpl::root(&scope);
+
+		let runtime = Arc::new(MockRuntime::default());
+		let graph = StagedDoubleNode::new(SourceNode(21.0f64), SourceNode(RuntimeHandle(runtime.clone())), SourceNode(8u64));
+
+		assert_eq!(GNode::eval(&graph, &ctx), GPoll::Pending);
+		assert_eq!(STAGED_RUNS.load(Ordering::Relaxed), 1, "the prologue runs synchronously on the miss");
+		assert_eq!(GNode::eval(&graph, &ctx), GPoll::Pending);
+		assert_eq!(STAGED_RUNS.load(Ordering::Relaxed), 1, "in flight must not rerun the prologue");
+		assert_eq!(runtime.drain(), vec![8]);
+		assert_eq!(GNode::eval(&graph, &ctx), GPoll::Final(42.0));
+		assert_eq!(STAGED_RUNS.load(Ordering::Relaxed), 1);
+	}
+
+	#[test]
+	fn prologue_interrupt_defers_the_spawn() {
+		let arena = Arena::new(64);
+		let generations = [];
+		let scope = scope_fixture(&generations, &arena);
+		let ctx = ContextImpl::root(&scope);
+
+		let gate = Arc::new(std::sync::atomic::AtomicBool::new(false));
+		let runtime = Arc::new(MockRuntime::default());
+		let graph = StagedSumNode::new(
+			SourceNode(40.0f64),
+			GatedSource(gate.clone(), 2.0),
+			SourceNode(RuntimeHandle(runtime.clone())),
+			SourceNode(9u64),
+		);
+
+		assert_eq!(GNode::eval(&graph, &ctx), GPoll::Pending);
+		assert_eq!(runtime.drain(), vec![], "an interrupted prologue must not spawn or claim the slot");
+		gate.store(true, Ordering::Relaxed);
+		assert_eq!(GNode::eval(&graph, &ctx), GPoll::Pending);
+		assert_eq!(runtime.drain(), vec![9]);
+		assert_eq!(GNode::eval(&graph, &ctx), GPoll::Final(42.0));
+	}
+
+	#[test]
+	fn async_kernels_read_captured_varargs() {
+		let arena = Arena::new(64);
+		let generations = [];
+		let scope = scope_fixture(&generations, &arena);
+		let root = ContextImpl::root(&scope);
+		let payload = 21.5f64;
+		let link = VarArgLink {
+			args: VarArgSlots::Single(&payload),
+			outer: None,
+		};
+		let ctx = root.with_varargs(&link);
+
+		let runtime = Arc::new(MockRuntime::default());
+		let graph = SnapshotVarargNode::new(SourceNode(()), SourceNode(RuntimeHandle(runtime.clone())), SourceNode(5u64));
+
+		assert_eq!(GNode::eval(&graph, &ctx), GPoll::Pending);
+		runtime.drain();
+		assert_eq!(GNode::eval(&graph, &ctx), GPoll::Final(21.5));
+	}
+
 	#[test]
 	fn async_kernels_read_the_captured_context_snapshot() {
 		let arena = Arena::new(64);
--- a/node-graph/node-macro/src/codegen.rs
+++ b/node-graph/node-macro/src/codegen.rs
@@ -111,8 +111,9 @@ pub(crate) fn generate_node_code(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
 		quote! { pub(super) #name: #r#gen }
 	});

+	let async_source = *is_async || crate::gcodegen::is_source_kernel(output_type);
 	let slot_value_type = crate::gcodegen::slot_value_type(output_type);
-	let slot_field = is_async
+	let slot_field = async_source
 		.then(|| quote! { pub(super) slot: std::sync::Arc<std::sync::Mutex<std::collections::HashMap<u64, Option<gcore::gpoll::GPoll<#slot_value_type>>>>> })
 		.into_iter();
 	let struct_fields = data_field_defs.chain(regular_field_defs).chain(slot_field);
@@ -223,11 +224,11 @@ pub(crate) fn generate_node_code(crate_ident: &CrateIdent, parsed: &ParsedNodeFn
 	let regular_inits = regular_field_names.iter().map(|name| {
 		quote! { #name }
 	});
-	let slot_init = is_async.then(|| quote! { slot: Default::default() }).into_iter();
+	let slot_init = async_source.then(|| quote! { slot: Default::default() }).into_iter();
 	let all_field_inits = data_inits.chain(regular_inits).chain(slot_init);

 	// Data fields may not implement Copy, PartialEq, etc., so only derive Debug and Clone
-	let struct_derives = if data_fields.is_empty() && !is_async {
+	let struct_derives = if data_fields.is_empty() && !async_source {
 		quote!(#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)])
 	} else {
 		quote!(#[derive(Debug, Clone)])
--- a/node-graph/node-macro/src/gcodegen.rs
+++ b/node-graph/node-macro/src/gcodegen.rs
@@ -18,14 +18,16 @@ pub(crate) fn generate_gnode_code(crate_ident: &CrateIdent, parsed: &ParsedNodeF
 		Some(ctx_param) => ctx_param.ident.clone(),
 		None => format_ident!("__Ctx"),
 	};
-	let async_source = parsed.is_async;
-	if async_source && parsed.fields.iter().any(|field| matches!(field.ty, ParsedFieldType::Node(_))) {
+	let async_fn = parsed.is_async;
+	let future_kernel = is_source_kernel(&parsed.output_type);
+	let async_source = async_fn || future_kernel;
+	if async_fn && parsed.fields.iter().any(|field| matches!(field.ty, ParsedFieldType::Node(_))) {
 		return Ok(GNodeTokens {
 			in_mod: quote!(),
 			top_level: quote!(),
 		});
 	}
-	let snapshot_ctx = async_source && matches!(&parsed.input.ty, Type::Path(path) if path.path.segments.last().is_some_and(|segment| segment.ident == "CtxSnapshot"));
+	let snapshot_ctx = async_fn && matches!(&parsed.input.ty, Type::Path(path) if path.path.segments.last().is_some_and(|segment| segment.ident == "CtxSnapshot"));

 	let mut ctx_bounds: Vec<TokenStream2> = match ctx_param {
 		Some(ctx_param) => ctx_param
@@ -80,11 +82,9 @@ pub(crate) fn generate_gnode_code(crate_ident: &CrateIdent, parsed: &ParsedNodeF
 	let mod_name = format_ident!("_{}_mod", parsed.mod_name);
 	let struct_name = format_ident!("{}Node", parsed.struct_name);
 	let output_type = &parsed.output_type;
-	let trait_output = match kernel_kind(&parsed.output_type) {
-		KernelKind::Interrupt(inner) | KernelKind::Poll(inner) => inner,
-		KernelKind::Plain => parsed.output_type.clone(),
-	};
+	let trait_output = slot_value_type(&parsed.output_type);
 	let raw_lazy = matches!(kernel_kind(&parsed.output_type), KernelKind::Poll(_));
+	let injected_name = |ident: &Ident| async_source && (ident == "_runtime" || ident == "_source");
 	let where_predicates: Vec<TokenStream2> = parsed.where_clause.iter().flat_map(|clause| clause.predicates.iter()).map(|predicate| quote!(#predicate)).collect();

 	let (data_fields, regular_fields): (Vec<_>, Vec<_>) = parsed.fields.iter().partition(|field| field.is_data_field);
@@ -121,7 +121,7 @@ pub(crate) fn generate_gnode_code(crate_ident: &CrateIdent, parsed: &ParsedNodeF
 		false => quote!(#core_types::gnode::GNode<#ctx_ident, Output = #output_type>),
 	};

-	let kernel_params = regular_fields.iter().map(|field| {
+	let kernel_params = regular_fields.iter().filter(|field| !injected_name(&field.pat_ident.ident)).map(|field| {
 		let pat = &field.pat_ident;
 		match &field.ty {
 			ParsedFieldType::Regular(RegularParsedField { ty, .. }) => quote!(#pat: #ty),
@@ -144,9 +144,10 @@ pub(crate) fn generate_gnode_code(crate_ident: &CrateIdent, parsed: &ParsedNodeF
 		}
 	});

-	let async_bounds = match async_source {
-		false => Vec::new(),
-		true => {
+	let async_bounds = match (async_fn, future_kernel) {
+		(false, false) => Vec::new(),
+		(false, true) => vec![quote!(#trait_output: Clone)],
+		(true, _) => {
 			let output_clone = std::iter::once(quote!(#trait_output: Clone));
 			let value_clones = regular_fields.iter().filter_map(|field| match &field.ty {
 				ParsedFieldType::Regular(RegularParsedField { ty, .. }) => Some(quote!(#ty: Clone)),
@@ -198,7 +199,7 @@ pub(crate) fn generate_gnode_code(crate_ident: &CrateIdent, parsed: &ParsedNodeF
 		(!tokens.is_empty()).then_some(tokens)
 	});

-	let call_args = regular_fields.iter().map(|field| {
+	let call_args = regular_fields.iter().filter(|field| !injected_name(&field.pat_ident.ident)).map(|field| {
 		let name = &field.pat_ident.ident;
 		match &field.ty {
 			ParsedFieldType::Node(_) if raw_lazy => quote!(&self.#name),
@@ -254,9 +255,8 @@ pub(crate) fn generate_gnode_code(crate_ident: &CrateIdent, parsed: &ParsedNodeF
 	let fn_where = &parsed.where_clause;
 	let body = &parsed.body;
 	let vis = &parsed.vis;
-	let injected = |field: &&&ParsedField| async_source && (field.pat_ident.ident == "_runtime" || field.pat_ident.ident == "_source");
-	let kernel_fields: Vec<&&ParsedField> = regular_fields.iter().filter(|field| !injected(field)).collect();
-	let kernel = match async_source {
+	let kernel_fields: Vec<&&ParsedField> = regular_fields.iter().filter(|field| !injected_name(&field.pat_ident.ident)).collect();
+	let kernel = match async_fn {
 		false => quote! {
 			#[allow(clippy::too_many_arguments)]
 			#vis fn #fn_name<#(#generics,)*>(#ctx_pat: &#ctx_ident #(, #data_params)* #(, #kernel_params)*) -> #output_type #fn_where #body
@@ -301,43 +301,52 @@ pub(crate) fn generate_gnode_code(crate_ident: &CrateIdent, parsed: &ParsedNodeF
 			}
 		},
 		KernelKind::Poll(_) => quote!(__cell.merge(#kernel_call)),
-		KernelKind::Plain => quote!(__cell.finish(#kernel_call)),
+		_ => quote!(__cell.finish(#kernel_call)),
 	};

-	let eval_tail = match async_source {
-		false => lift,
-		true => {
+	let placeholder_value_names: Vec<&Ident> = kernel_fields
+		.iter()
+		.filter(|field| matches!(field.ty, ParsedFieldType::Regular(_)))
+		.map(|field| &field.pat_ident.ident)
+		.collect();
+	let inflight = match &parsed.attributes.placeholder {
+		Some(path) => quote!(__cell.merge(#core_types::gpoll::GPoll::Partial(#path(#(&#placeholder_value_names),*)))),
+		None => quote!(#core_types::gpoll::GPoll::Pending),
+	};
+	let slot_check = quote! {
+		let __key = #core_types::wire::cache_key(__input);
+		{
+			let __entries = self.slot.lock().unwrap();
+			if let Some(__state) = __entries.get(&__key) {
+				return match __state {
+					Some(value) => __cell.merge(value.clone()),
+					None => #inflight,
+				};
+			}
+		}
+	};
+	let future_completion = |payload: &Type| match kernel_kind(payload) {
+		KernelKind::Poll(_) => quote!(__future.await),
+		KernelKind::Interrupt(_) => quote! {
+			match __future.await {
+				Ok(value) => #core_types::gpoll::GPoll::Final(value),
+				Err(interrupt) => interrupt.into(),
+			}
+		},
+		_ => quote!(#core_types::gpoll::GPoll::Final(__future.await)),
+	};
+	let eval_tail = match (async_fn, future_kernel) {
+		(false, false) => lift,
+		(true, _) => {
 			let kernel_value_names: Vec<&Ident> = kernel_fields.iter().map(|field| &field.pat_ident.ident).collect();
-			let inflight = match &parsed.attributes.placeholder {
-				Some(path) => quote!(__cell.merge(#core_types::gpoll::GPoll::Partial(#path(#(&#kernel_value_names),*)))),
-				None => quote!(#core_types::gpoll::GPoll::Pending),
-			};
 			let snapshot_binding = snapshot_ctx.then(|| quote!(let __snapshot = #core_types::context::CtxSnapshot::capture(__input);)).into_iter();
 			let snapshot_arg = snapshot_ctx.then(|| quote!(__snapshot)).into_iter();
 			let future_args = snapshot_arg
 				.chain(data_names.iter().map(|name| quote!(self.#name.clone())))
 				.chain(kernel_value_names.iter().map(|name| quote!(#name.clone())));
-			let completion = match kernel_kind(&parsed.output_type) {
-				KernelKind::Plain => quote!(#core_types::gpoll::GPoll::Final(__future.await)),
-				KernelKind::Poll(_) => quote!(__future.await),
-				KernelKind::Interrupt(_) => quote! {
-					match __future.await {
-						Ok(value) => #core_types::gpoll::GPoll::Final(value),
-						Err(interrupt) => interrupt.into(),
-					}
-				},
-			};
+			let completion = future_completion(&parsed.output_type);
 			quote! {
-				let __key = #core_types::wire::cache_key(__input);
-				{
-					let __entries = self.slot.lock().unwrap();
-					if let Some(__state) = __entries.get(&__key) {
-						return match __state {
-							Some(value) => __cell.merge(value.clone()),
-							None => #inflight,
-						};
-					}
-				}
+				#slot_check
 				self.slot.lock().unwrap().insert(__key, None);
 				let __slot = std::sync::Arc::clone(&self.slot);
 				#(#snapshot_binding)*
@@ -349,6 +358,41 @@ pub(crate) fn generate_gnode_code(crate_ident: &CrateIdent, parsed: &ParsedNodeF
 				#inflight
 			}
 		}
+		(false, true) => {
+			let (placeholder_binding, spawn_return) = match &parsed.attributes.placeholder {
+				Some(path) => (
+					quote!(let __placeholder = #path(#(&#placeholder_value_names),*);),
+					quote!(__cell.merge(#core_types::gpoll::GPoll::Partial(__placeholder))),
+				),
+				None => (quote!(), quote!(#core_types::gpoll::GPoll::Pending)),
+			};
+			let acquire = match kernel_kind(&parsed.output_type) {
+				KernelKind::FutureInterrupt(_) => quote! {
+					let __future = match #kernel_call {
+						Ok(future) => future,
+						Err(interrupt) => return interrupt.into(),
+					};
+				},
+				_ => quote!(let __future = #kernel_call;),
+			};
+			let payload = match kernel_kind(&parsed.output_type) {
+				KernelKind::Future(payload) | KernelKind::FutureInterrupt(payload) => payload,
+				_ => unreachable!("guarded by future_kernel"),
+			};
+			let completion = future_completion(&payload);
+			quote! {
+				#slot_check
+				#placeholder_binding
+				#acquire
+				self.slot.lock().unwrap().insert(__key, None);
+				let __slot = std::sync::Arc::clone(&self.slot);
+				_runtime.0.spawn(_source, Box::pin(async move {
+					let __value = #completion;
+					__slot.lock().unwrap().insert(__key, Some(__value));
+				}));
+				#spawn_return
+			}
+		}
 	};

 	let wire = entries_tokens(parsed, &struct_name, &data_field_generic_idents, &regular_fields);
@@ -406,13 +450,36 @@ pub(crate) fn slot_value_type(output: &Type) -> Type {
 	match kernel_kind(output) {
 		KernelKind::Plain => output.clone(),
 		KernelKind::Poll(inner) | KernelKind::Interrupt(inner) => inner,
+		KernelKind::Future(payload) | KernelKind::FutureInterrupt(payload) => match kernel_kind(&payload) {
+			KernelKind::Poll(inner) | KernelKind::Interrupt(inner) => inner,
+			_ => payload,
+		},
 	}
 }

+pub(crate) fn is_source_kernel(output: &Type) -> bool {
+	matches!(kernel_kind(output), KernelKind::Future(_) | KernelKind::FutureInterrupt(_))
+}
+
 enum KernelKind {
 	Plain,
 	Interrupt(Type),
 	Poll(Type),
+	Future(Type),
+	FutureInterrupt(Type),
+}
+
+fn source_future_payload(segment: &syn::PathSegment) -> Type {
+	let PathArguments::AngleBracketed(args) = &segment.arguments else {
+		return syn::parse_quote!(());
+	};
+	args.args
+		.iter()
+		.find_map(|argument| match argument {
+			GenericArgument::Type(ty) => Some(ty.clone()),
+			_ => None,
+		})
+		.unwrap_or_else(|| syn::parse_quote!(()))
 }

 fn kernel_kind(output: &Type) -> KernelKind {
@@ -428,6 +495,7 @@ fn kernel_kind(output: &Type) -> KernelKind {
 			});
 			inner.map(KernelKind::Poll).unwrap_or_else(plain)
 		}
+		"SourceFuture" => KernelKind::Future(source_future_payload(segment)),
 		"Result" => {
 			let PathArguments::AngleBracketed(args) = &segment.arguments else { return plain() };
 			let mut types = args.args.iter().filter_map(|argument| match argument {
@@ -437,10 +505,17 @@ fn kernel_kind(output: &Type) -> KernelKind {
 			let (Some(inner), Some(Type::Path(error_path))) = (types.next(), types.next()) else {
 				return plain();
 			};
-			match error_path.path.segments.last().is_some_and(|segment| segment.ident == "Interrupt") {
-				true => KernelKind::Interrupt(inner.clone()),
-				false => plain(),
+			if !error_path.path.segments.last().is_some_and(|segment| segment.ident == "Interrupt") {
+				return plain();
 			}
+			if let Type::Path(inner_path) = inner {
+				if let Some(inner_segment) = inner_path.path.segments.last() {
+					if inner_segment.ident == "SourceFuture" {
+						return KernelKind::FutureInterrupt(source_future_payload(inner_segment));
+					}
+				}
+			}
+			KernelKind::Interrupt(inner.clone())
 		}
 		_ => plain(),
 	}
--- a/node-graph/node-macro/src/parsing.rs
+++ b/node-graph/node-macro/src/parsing.rs
@@ -997,7 +997,7 @@ pub fn new_node_fn(attr: TokenStream2, item: TokenStream2) -> syn::Result<TokenS
 	let crate_ident = CrateIdent::default();
 	let mut parsed_node = parse_node_fn(attr, item.clone()).map_err(|e| Error::new(e.span(), format!("Failed to parse node function:\n{e}")))?;
 	parsed_node.replace_impl_trait_in_input();
-	if parsed_node.is_async {
+	if parsed_node.is_async || crate::gcodegen::is_source_kernel(&parsed_node.output_type) {
 		let core_types = crate_ident.gcore()?.clone();
 		parsed_node.inject_async_source_fields(&core_types);
 	}
--- a/node-graph/node-macro/src/validation.rs
+++ b/node-graph/node-macro/src/validation.rs
@@ -23,18 +23,30 @@ pub fn validate_node_fn(parsed: &ParsedNodeFn) -> syn::Result<()> {

 fn validate_async_source(parsed: &ParsedNodeFn) {
 	let snapshot_ctx = matches!(&parsed.input.ty, Type::Path(path) if path.path.segments.last().is_some_and(|segment| segment.ident == "CtxSnapshot"));
+	let future_kernel = crate::gcodegen::is_source_kernel(&parsed.output_type);
+	if parsed.is_async && future_kernel {
+		emit_error!(
+			parsed.output_type.span(),
+			"an `async fn` kernel already is the async part; returning `SourceFuture` is the sync-prologue form, so drop the `async` keyword or return the value directly"
+		);
+		return;
+	}
 	if !parsed.is_async {
 		if snapshot_ctx {
 			emit_error!(parsed.input.pat_ident.span(), "`CtxSnapshot` is the async source context; synchronous nodes take `impl Ctx` and read through extract bounds");
 		}
-		return;
+		if !future_kernel {
+			return;
+		}
 	}
-	for field in &parsed.fields {
-		if matches!(field.ty, ParsedFieldType::Node(_)) {
-			emit_error!(
-				field.pat_ident.span(),
-				"async source nodes cannot take `impl Node` inputs: the spawned future outlives any borrow of the graph, so it cannot evaluate other nodes; declare the input as an eager value instead"
-			);
+	if parsed.is_async {
+		for field in &parsed.fields {
+			if matches!(field.ty, ParsedFieldType::Node(_)) {
+				emit_error!(
+					field.pat_ident.span(),
+					"`async fn` source nodes cannot take `impl Node` inputs: the spawned future outlives any borrow of the graph, so it cannot evaluate other nodes; use the sync-prologue form (return `SourceFuture`) to evaluate lazy inputs before spawning"
+				);
+			}
 		}
 	}
 	let ctx_ident = match &parsed.input.ty {