diff --git a/dagx-macros/src/lib.rs b/dagx-macros/src/lib.rs index 76c1f1f..1b9006c 100644 --- a/dagx-macros/src/lib.rs +++ b/dagx-macros/src/lib.rs @@ -323,6 +323,112 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream { .iter() .any(|arg| matches!(arg, FnArg::Receiver(_))); + // Generate extract_and_run implementation based on parameter count + let extract_and_run_impl = match params.len() { + 0 => { + // Zero parameters - no extraction needed + quote! { + fn extract_and_run( + self, + _receivers: Vec>, + ) -> impl std::future::Future> + Send { + async move { + let input = (); + Ok(self.run(input).await) + } + } + } + } + 1 => { + // Single parameter + let param_type = ¶ms[0].1; + quote! { + fn extract_and_run( + self, + mut receivers: Vec>, + ) -> impl std::future::Future> + Send { + async move { + use futures::channel::oneshot; + use std::sync::Arc; + + if receivers.len() != 1 { + return Err(format!("Expected 1 dependency, got {}", receivers.len())); + } + + let rx = *receivers.pop() + .unwrap() + .downcast::>>() + .map_err(|_| format!("Type mismatch: expected Arc<{}>", std::any::type_name::<#param_type>()))?; + + let arc_value = rx.await + .map_err(|_| "Channel closed before receiving value".to_string())?; + + let input = (*arc_value).clone(); + Ok(self.run(input).await) + } + } + } + } + _ => { + // Multiple parameters + let param_count = params.len(); + let param_types: Vec<_> = params.iter().map(|(_, ty)| ty).collect(); + let indices: Vec<_> = (0..param_count).collect(); + + // Generate unique receiver variable names + let rx_vars: Vec<_> = (0..param_count) + .map(|i| syn::Ident::new(&format!("rx_{}", i), proc_macro2::Span::call_site())) + .collect(); + + // Create syn::Index for tuple field access (avoids the suffix warning) + let syn_indices: Vec<_> = (0..param_count).map(syn::Index::from).collect(); + + quote! { + fn extract_and_run( + self, + receivers: Vec>, + ) -> impl std::future::Future> + Send { + async move { + use futures::channel::oneshot; + use std::sync::Arc; + + let expected_count = #param_count; + if receivers.len() != expected_count { + return Err(format!("Expected {} dependencies, got {}", expected_count, receivers.len())); + } + + let mut iter = receivers.into_iter(); + + // Extract each receiver + #( + let #rx_vars = *iter.next() + .ok_or_else(|| format!("Missing receiver at index {}", #indices))? + .downcast::>>() + .map_err(|_| format!("Type mismatch at index {}: expected Arc<{}>", + #indices, std::any::type_name::<#param_types>()))?; + )* + + // Await all channels concurrently + let arc_results = futures::join!( + #( + async move { + #rx_vars.await.map_err(|_| format!("Channel {} closed", #indices)) + } + ),* + ); + + // Clone inner values and build tuple + let input = (#( + (*arc_results.#syn_indices?).clone() + ),*); + + Ok(self.run(input).await) + } + } + } + } + }; + // Generate the Task trait implementation based on whether we have self and async/sync let expanded = if has_self_receiver { // Stateful task - consumes self but delegates to a method that borrows @@ -337,6 +443,8 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream { let #param_destructure = input; self.run_impl(#param_refs).await } + + #extract_and_run_impl } impl #struct_name { @@ -354,6 +462,8 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream { let #param_destructure = input; self.run_impl(#param_refs) } + + #extract_and_run_impl } impl #struct_name { @@ -374,6 +484,8 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream { let #param_destructure = input; Self::run_impl(#param_refs).await } + + #extract_and_run_impl } impl #struct_name { @@ -392,6 +504,8 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream { let #param_destructure = input; Self::run_impl(#param_refs) } + + #extract_and_run_impl } impl #struct_name { diff --git a/src/extract.rs b/src/extract.rs index 83fe7e9..a926b45 100644 --- a/src/extract.rs +++ b/src/extract.rs @@ -28,6 +28,11 @@ use futures::channel::oneshot; /// /// Type erasure occurs only at the ExecutableNode trait boundary - by the time /// we're in ExtractInput, we know the concrete type and can safely downcast. +/// +/// NOTE: This trait is now primarily used by the internal `task_fn` test helper. +/// The `#[task]` macro generates inline extraction logic via `extract_and_run()`, +/// which allows ANY type to work without requiring ExtractInput implementations. +#[allow(dead_code)] pub(crate) trait ExtractInput: Sized + Clone { fn extract_from_channels( receivers: Vec>, diff --git a/src/node.rs b/src/node.rs index 0d2c44e..ee0a24c 100644 --- a/src/node.rs +++ b/src/node.rs @@ -14,7 +14,6 @@ use std::sync::Arc; use futures::channel::oneshot; use crate::error::DagError; -use crate::extract::ExtractInput; use crate::task::Task; use crate::types::NodeId; @@ -84,7 +83,7 @@ impl TypedNode { impl ExecutableNode for TypedNode where - T::Input: 'static + Clone + ExtractInput, + T::Input: 'static + Clone, T::Output: 'static + Clone, { fn create_output_channels(&self, num_dependents: usize) -> ChannelPair { @@ -107,24 +106,22 @@ where senders: ErasedChannels, ) -> ExecuteFuture { Box::pin(async move { - // Extract typed input from type-erased receivers - // TypedNode knows T::Input, so it can safely downcast - let input = match T::Input::extract_from_channels(receivers).await { - Ok(input) => input, + // Use the task's extract_and_run method which has inline extraction logic + // generated by the #[task] macro. This allows ANY type to work without + // requiring ExtractInput trait implementations. + let output = match self.task.extract_and_run(receivers).await { + Ok(output) => output, Err(_msg) => { // Failed to extract input - likely channel closed or type mismatch return Err(DagError::InvalidDependency { task_id: self.id.0 }); } }; - // Execute task (consumes self.task) - let output = self.task.run(input).await; - // Arc Wrapping Strategy: // Wrap output in Arc ONCE for efficient fan-out sharing. // For 1:N dependencies, this enables O(1) sharing (Arc clone = atomic increment) // instead of O(N) cloning (N deep copies of output data). - // Downstream tasks receive &T after ExtractInput extracts from Arc. + // Downstream tasks receive &T after extraction from Arc. let arc_output = Arc::new(output); // Send Arc to all dependents via their channels diff --git a/src/runner.rs b/src/runner.rs index ec4120d..f75aa35 100644 --- a/src/runner.rs +++ b/src/runner.rs @@ -22,7 +22,6 @@ use tracing::{debug, error, info, trace}; use crate::builder::TaskBuilder; use crate::error::{DagError, DagResult}; -use crate::extract::ExtractInput; use crate::node::{ExecutableNode, TypedNode}; use crate::task::Task; use crate::types::{NodeId, Pending, TaskHandle}; @@ -196,7 +195,7 @@ impl DagRunner { pub fn add_task(&self, task: Tk) -> TaskBuilder<'_, Tk, Pending> where Tk: Task + 'static, - Tk::Input: 'static + Clone + ExtractInput, + Tk::Input: 'static + Clone, Tk::Output: 'static + Clone, { let id = self.alloc_id(); diff --git a/src/task.rs b/src/task.rs index a25f8cd..f569689 100644 --- a/src/task.rs +++ b/src/task.rs @@ -93,6 +93,17 @@ pub trait Task: Send { type Output: Send + Sync; // Sync required for Arc-wrapping and cross-thread sharing fn run(self, input: Self::Input) -> impl Future + Send; + + /// Internal method: Extract input from type-erased channels and execute the task. + /// + /// This method is implemented automatically by the #[task] macro with inline + /// extraction logic based on the task's parameter signature. This allows ANY + /// type to work without requiring trait implementations. + #[doc(hidden)] + fn extract_and_run( + self, + receivers: Vec>, + ) -> impl Future> + Send; } /// Convenience function to create a task from a closure. @@ -104,7 +115,7 @@ pub fn task_fn(f: F) -> impl Task where F: FnMut(I) -> Fut + Send + 'static, Fut: Future + Send + 'static, - I: Send + 'static, + I: Send + crate::extract::ExtractInput + 'static, O: Send + Sync + 'static, { struct TaskFn @@ -122,7 +133,7 @@ where where F: FnMut(I) -> Fut + Send, Fut: Future + Send, - I: Send, + I: Send + crate::extract::ExtractInput, O: Send + Sync, { type Input = I; @@ -131,6 +142,18 @@ where async fn run(mut self, input: I) -> O { (self.f)(input).await } + + async fn extract_and_run( + self, + receivers: Vec>, + ) -> Result { + // For task_fn, we use ExtractInput since we don't have macro-generated extraction + let input = I::extract_from_channels(receivers) + .await + .map_err(|e| format!("Failed to extract input: {}", e))?; + + Ok(self.run(input).await) + } } TaskFn { diff --git a/src/task/tests.rs b/src/task/tests.rs index 2f60539..ee2b74b 100644 --- a/src/task/tests.rs +++ b/src/task/tests.rs @@ -83,6 +83,17 @@ impl Task for CustomTask { self.state }) } + + #[allow(refining_impl_trait)] + fn extract_and_run( + self, + _receivers: Vec>, + ) -> Pin> + Send>> { + Box::pin(async move { + // This is a test helper - just return an error + Err("extract_and_run not implemented for CustomTask (test helper)".to_string()) + }) + } } #[tokio::test] diff --git a/tests/custom_type_test.rs b/tests/custom_type_test.rs new file mode 100644 index 0000000..623d330 --- /dev/null +++ b/tests/custom_type_test.rs @@ -0,0 +1,137 @@ +// Test that custom types work without any trait implementations +use dagx::{task, DagRunner, Task}; + +#[derive(Clone, Debug, PartialEq)] +struct Person { + name: String, + age: u32, +} + +#[derive(Clone, Debug, PartialEq)] +struct Company { + name: String, + employees: Vec, +} + +// Source task returning custom type +struct FetchPerson; + +#[task] +impl FetchPerson { + async fn run(&self) -> Person { + Person { + name: "Alice".to_string(), + age: 30, + } + } +} + +// Task consuming custom type +struct ProcessPerson; + +#[task] +impl ProcessPerson { + async fn run(person: &Person) -> String { + format!("{} is {} years old", person.name, person.age) + } +} + +// Task transforming custom type to another custom type +struct PersonToCompany; + +#[task] +impl PersonToCompany { + async fn run(person: &Person) -> Company { + Company { + name: format!("{}'s Company", person.name), + employees: vec![person.clone()], + } + } +} + +#[tokio::test] +async fn test_custom_type_single_dependency() { + let dag = DagRunner::new(); + + let fetch = dag.add_task(FetchPerson); + let process = dag.add_task(ProcessPerson).depends_on(&fetch); + + dag.run(|fut| { + tokio::spawn(fut); + }) + .await + .unwrap(); + + let result = dag.get(process).unwrap(); + assert_eq!(result, "Alice is 30 years old"); +} + +#[tokio::test] +async fn test_custom_type_passthrough() { + let dag = DagRunner::new(); + + let fetch = dag.add_task(FetchPerson); + let transform = dag.add_task(PersonToCompany).depends_on(&fetch); + let process = dag.add_task(ProcessPerson).depends_on(&fetch); + + dag.run(|fut| { + tokio::spawn(fut); + }) + .await + .unwrap(); + + let person_result = dag.get(process).unwrap(); + let company_result = dag.get(transform).unwrap(); + + assert_eq!(person_result, "Alice is 30 years old"); + assert_eq!(company_result.name, "Alice's Company"); + assert_eq!(company_result.employees.len(), 1); +} + +#[tokio::test] +async fn test_custom_type_diamond_pattern() { + let dag = DagRunner::new(); + + // Source + let fetch = dag.add_task(FetchPerson); + + // Two paths from source + let path_a = dag.add_task(ProcessPerson).depends_on(&fetch); + let path_b = dag.add_task(PersonToCompany).depends_on(&fetch); + + dag.run(|fut| { + tokio::spawn(fut); + }) + .await + .unwrap(); + + let result_a = dag.get(path_a).unwrap(); + let result_b = dag.get(path_b).unwrap(); + + assert_eq!(result_a, "Alice is 30 years old"); + assert_eq!(result_b.name, "Alice's Company"); +} + +#[tokio::test] +async fn test_custom_type_fan_out() { + let dag = DagRunner::new(); + + let fetch = dag.add_task(FetchPerson); + let process1 = dag.add_task(ProcessPerson).depends_on(&fetch); + let process2 = dag.add_task(ProcessPerson).depends_on(&fetch); + let process3 = dag.add_task(ProcessPerson).depends_on(&fetch); + + dag.run(|fut| { + tokio::spawn(fut); + }) + .await + .unwrap(); + + let result1 = dag.get(process1).unwrap(); + let result2 = dag.get(process2).unwrap(); + let result3 = dag.get(process3).unwrap(); + + assert_eq!(result1, "Alice is 30 years old"); + assert_eq!(result2, "Alice is 30 years old"); + assert_eq!(result3, "Alice is 30 years old"); +}