diff --git a/node-graph/libraries/core-types/src/runtime.rs b/node-graph/libraries/core-types/src/runtime.rs index e12cb77a90..0d42fc82a4 100644 --- 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 + Send + 'static>>; +pub type SourceFuture = Pin + Send + 'static>>; #[cfg(target_family = "wasm")] -pub type SourceFuture = Pin + 'static>>; +pub type SourceFuture = Pin + '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::()).copied().unwrap_or(0.) + } + + static STAGED_RUNS: AtomicU32 = AtomicU32::new(0); + + #[node_macro::node(category(""))] + fn staged_double(_: impl Ctx, value: f64) -> SourceFuture { + 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, Output = f64>) -> Result, crate::gpoll::Interrupt> { + let addend = addend.eval(ctx)?; + Ok(Box::pin(async move { value + addend })) + } + + struct GatedSource(Arc, f64); + + impl GNode for GatedSource { + type Output = f64; + + fn eval(&self, _input: &Input) -> GPoll { + 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); diff --git a/node-graph/node-macro/src/codegen.rs b/node-graph/node-macro/src/codegen.rs index d48a4bf6e2..045e63d6f3 100644 --- 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>>>> }) .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)]) diff --git a/node-graph/node-macro/src/gcodegen.rs b/node-graph/node-macro/src/gcodegen.rs index 0c20febdd0..693ad783fc 100644 --- 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 = 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 = 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, ®ular_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(), } diff --git a/node-graph/node-macro/src/parsing.rs b/node-graph/node-macro/src/parsing.rs index 213b84e861..405d08f784 100644 --- 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 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 {