swaits/dagx
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove clone bound from task data

Browse Source
This commit is contained in:
TechnoPorg committed 2026-02-11 21:59:41 +01:00
1 parent 064c2a15be
commit 0137296fb4
9 files changed
+233 -442
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

dagx-macros/src/lib.rs
+49 -196
View File
@@ -4,8 +4,9 @@
//! the `Task` trait by deriving Input and Output types from the `run()` method signature.
use proc_macro::TokenStream;
use proc_macro2::Span;
use quote::quote;
use syn::{parse_macro_input, FnArg, ImplItem, ItemImpl, Pat, PatType, ReturnType, Type};
use syn::{parse_macro_input, FnArg, Ident, ImplItem, ItemImpl, Pat, PatType, ReturnType, Type};
/// Attribute macro to automatically implement the `Task` trait.
///
@@ -18,8 +19,7 @@ use syn::{parse_macro_input, FnArg, ImplItem, ItemImpl, Pat, PatType, ReturnType
/// - Handles various input patterns (no inputs, single input, multiple inputs)
///
/// **Key Feature**: Custom types work automatically without implementing any traits!
/// The macro generates inline extraction logic in `extract_and_run()` specific to your
/// task's parameter types. Just derive `Clone` on your types and they'll work seamlessly.
/// The macro generates inline extraction logic specific to your task's parameter types.
///
/// # Task Patterns
///
@@ -135,8 +135,6 @@ use syn::{parse_macro_input, FnArg, ImplItem, ItemImpl, Pat, PatType, ReturnType
/// - The impl block must contain exactly one `async fn run()` method
/// - The `run()` method can be stateless (no self parameter) or stateful (`&mut self`)
/// - All input parameters must be references (e.g., `&i32`, not `i32`)
/// - The macro requires `Task` to be in scope: `use dagx::Task;`
/// - For stateless tasks, the struct must implement `Default` (e.g., unit structs)
///
/// # Generated Code
///
@@ -232,9 +230,6 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream {
}
};
// Check if the method is async or sync
let is_async = run_method.sig.asyncness.is_some();
// Extract parameters (excluding self)
let params_result: Result<Vec<_>, _> = run_method
.sig
@@ -280,32 +275,39 @@ pub fn task(_attr: TokenStream, item: TokenStream) -> TokenStream {
// Extract return type
let output_type = match &run_method.sig.output {
ReturnType::Default => syn::parse_quote!(()),
ReturnType::Type(_, ty) => ty.clone(),
ReturnType::Type(_, ty) => {
if let Type::Tuple(_tuple) = &**ty {
return syn::Error::new_spanned(
ty,
"Returning bare tuples from task functions is currently not supported.\n\n\
Please wrap this in a newtype or struct.",
)
.into_compile_error()
.into();
}
ty.clone()
}
};
// Build Input type based on parameter count
let input_type = match params.len() {
0 => quote! { () },
1 => {
let (_name, ty) = &params[0];
quote! { #ty }
}
_ => {
let types: Vec<_> = params.iter().map(|(_, ty)| ty).collect();
quote! { ( #(#types),* ) }
}
};
let types: Vec<_> = params.iter().map(|(_, ty)| ty).collect();
let input_type = quote! { ( #(#types,)* ) };
// Generate parameter destructuring for the wrapper run() method
let (param_destructure, param_refs) = if params.is_empty() {
(quote! { _ }, quote! {})
} else if params.len() == 1 {
let (name, _) = &params[0];
(quote! { #name }, quote! { &#name })
} else {
let names: Vec<_> = params.iter().map(|(name, _)| name).collect();
let refs: Vec<_> = params.iter().map(|(name, _)| quote! { &#name }).collect();
(quote! { ( #(#names),* ) }, quote! { #(#refs),* })
// Generate parameter references for the wrapper run() method
let (param_destructure, param_refs) = {
let (destructure, refs): (Vec<_>, Vec<_>) = params
.iter()
.map(|(ident, _ty)| {
let ident = Ident::new(&format!("__dagx_extract_{}", ident), Span::mixed_site());
(
quote! {
let (#ident, input) = input.next();
},
ident,
)
})
.unzip();
(quote! { #(#destructure)* }, quote! { #(#refs),* })
};
// Clone the run method and rename it to run_impl
@@ -319,179 +321,30 @@ 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,
_dependencies: Vec<std::sync::Arc<dyn std::any::Any + Send + Sync>>,
) -> impl std::future::Future<Output = Result<Self::Output, String>> + Send {
async move {
let input = ();
Ok(self.run(input).await)
}
}
}
}
1 => {
// Single parameter
let param_type = &params[0].1;
quote! {
fn extract_and_run(
self,
mut dependencies: Vec<std::sync::Arc<dyn std::any::Any + Send + Sync>>,
) -> impl std::future::Future<Output = Result<Self::Output, String>> + Send {
async move {
use std::sync::Arc;
let is_async = run_method.sig.asyncness.is_some();
if dependencies.len() != 1 {
return Err(format!("Expected 1 dependency, got {}", dependencies.len()));
}
let arc_value = dependencies.pop()
.unwrap()
.downcast::<#param_type>()
.map_err(|_| format!("Type mismatch: expected Arc<{}>", std::any::type_name::<#param_type>()))?;
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();
// 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,
dependencies: Vec<std::sync::Arc<dyn std::any::Any + Send + Sync>>,
) -> impl std::future::Future<Output = Result<Self::Output, String>> + Send {
async move {
use std::sync::Arc;
let expected_count = #param_count;
if dependencies.len() != expected_count {
return Err(format!("Expected {} dependencies, got {}", expected_count, dependencies.len()));
}
let mut iter = dependencies.into_iter();
// Extract each parameter
let arc_results = (#(
iter.next()
.ok_or_else(|| format!("Missing parameter at index {}", #indices))?
.downcast::<#param_types>()
.map_err(|_| format!("Type mismatch at index {}: expected Arc<{}>",
#indices, std::any::type_name::<#param_types>()))?,
)*);
// Clone inner values and build tuple
let input = (#(
(*arc_results.#syn_indices).clone()
),*);
Ok(self.run(input).await)
}
}
}
}
let run_call = match (has_self_receiver, is_async) {
(true, true) => quote! { self.run_impl(#param_refs).await },
(true, false) => quote! { self.run_impl(#param_refs) },
(false, true) => quote! { Self::run_impl(#param_refs).await },
(false, false) => quote! { Self::run_impl(#param_refs) },
};
// 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
if is_async {
// Async with self
quote! {
impl ::dagx::Task for #struct_name {
type Input = #input_type;
type Output = #output_type;
quote! {
impl ::dagx::Task for #struct_name {
type Input = #input_type;
type Output = #output_type;
async fn run(mut self, input: Self::Input) -> Self::Output {
let #param_destructure = input;
self.run_impl(#param_refs).await
}
#extract_and_run_impl
}
impl #struct_name {
#run_impl_method
}
}
} else {
// Sync with self - wrap in async block
quote! {
impl ::dagx::Task for #struct_name {
type Input = #input_type;
type Output = #output_type;
async fn run(mut self, input: Self::Input) -> Self::Output {
let #param_destructure = input;
self.run_impl(#param_refs)
}
#extract_and_run_impl
}
impl #struct_name {
#run_impl_method
}
async fn run(mut self, mut input: ::dagx::TaskInput<'_, Self::Input>) -> Self::Output {
#param_destructure
#run_call
}
}
} else {
// Stateless task
if is_async {
// Async stateless
quote! {
impl ::dagx::Task for #struct_name {
type Input = #input_type;
type Output = #output_type;
async fn run(self, input: Self::Input) -> Self::Output {
let #param_destructure = input;
Self::run_impl(#param_refs).await
}
#extract_and_run_impl
}
impl #struct_name {
#[inline]
#run_impl_method
}
}
} else {
// Sync stateless
quote! {
impl ::dagx::Task for #struct_name {
type Input = #input_type;
type Output = #output_type;
async fn run(self, input: Self::Input) -> Self::Output {
let #param_destructure = input;
Self::run_impl(#param_refs)
}
#extract_and_run_impl
}
impl #struct_name {
#[inline]
#run_impl_method
}
}
impl #struct_name {
#run_impl_method
}
};
TokenStream::from(expanded)
}
.into()
}
src/builder.rs
-9
View File
@@ -11,15 +11,6 @@ use crate::types::{NodeId, TaskHandle};
#[cfg(feature = "tracing")]
use tracing::debug;
/// Marker trait to identify the unit type at compile time.
///
/// This is used to enforce that tasks with non-unit inputs MUST call `.depends_on()`,
/// while tasks with unit inputs can be used directly.
pub trait IsUnitType {}
/// Only the unit type implements IsUnitType
impl IsUnitType for () {}
/// Node builder that tracks dependency completion via type state.
///
/// A `TaskBuilder<Tk, Deps>` is returned from [`DagRunner::add_task`] and uses the type-state
src/deps.rs
+1 -1
View File
@@ -50,7 +50,7 @@ impl DepsTuple<()> for () {
}
// Implementation for single dependency (all forms)
impl<T> DepsTuple<T::Output> for T
impl<T> DepsTuple<(T::Output,)> for T
where
T: IntoNodeRef,
{
src/extract.rs
+76 -146
View File
@@ -32,10 +32,11 @@
//!
//! This is standard practice in Rust for working with tuples of different arities.
use std::any::Any;
use std::sync::Arc;
use std::{collections::HashMap, hash::Hash, sync::Arc};
/// Helper trait for async input extraction from list of dependencies.
use crate::TaskInput;
/// Helper trait for input extraction from list of dependencies.
///
/// 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.
@@ -44,15 +45,21 @@ use std::sync::Arc;
/// 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_deps(
deps: Vec<Arc<dyn Any + Send + Sync>>,
) -> impl std::future::Future<Output = Result<Self, String>> + Send;
pub trait ExtractInput: Sized + Clone {
type Input: Send;
type Retv<'input>;
fn extract_from_task_input<'input>(
input: TaskInput<'input, Self::Input>,
) -> Result<Self::Retv<'input>, String>;
}
// Unit type - no dependencies
impl ExtractInput for () {
async fn extract_from_deps(_deps: Vec<Arc<dyn Any + Send + Sync>>) -> Result<Self, String> {
type Input = ();
type Retv<'input> = ();
fn extract_from_task_input(_input: TaskInput<'_, Self>) -> Result<Self, String> {
Ok(())
}
}
@@ -65,15 +72,10 @@ macro_rules! impl_extract_single {
($($t:ty),+) => {
$(
impl ExtractInput for $t {
async fn extract_from_deps(mut dependencies: Vec<Arc<dyn Any + Send + Sync>>) -> Result<Self, String> {
if dependencies.len() != 1 {
return Err(format!("Expected 1 dependency, got {}", dependencies.len()));
}
let arc_value = dependencies.pop()
.unwrap()
.downcast::<$t>()
.map_err(|_| format!("Type mismatch: expected Arc<{}>", std::any::type_name::<$t>()))?;
Ok((*arc_value).clone())
type Input = ($t,);
type Retv<'input> = &'input $t;
fn extract_from_task_input<'input>(input: TaskInput<'input, Self::Input>) -> Result<Self::Retv<'input>, String> {
Ok(input.next().0)
}
}
)+
@@ -90,24 +92,13 @@ where
T: 'static + Clone + Send + Sync,
E: 'static + Clone + Send + Sync,
{
async fn extract_from_deps(
mut dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> Result<Self, String> {
if dependencies.len() != 1 {
return Err(format!("Expected 1 dependency, got {}", dependencies.len()));
}
let arc_value = dependencies
.pop()
.unwrap()
.downcast::<Result<T, E>>()
.map_err(|_| {
format!(
"Type mismatch: expected Arc<Result<{}, {}>>",
std::any::type_name::<T>(),
std::any::type_name::<E>()
)
})?;
Ok((*arc_value).clone())
type Input = (Self,);
type Retv<'input> = &'input Result<T, E>;
fn extract_from_task_input<'input>(
input: TaskInput<'input, (Self,)>,
) -> Result<Self::Retv<'input>, String> {
Ok(input.next().0)
}
}
@@ -116,23 +107,13 @@ impl<T> ExtractInput for Option<T>
where
T: 'static + Clone + Send + Sync,
{
async fn extract_from_deps(
mut dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> Result<Self, String> {
if dependencies.len() != 1 {
return Err(format!("Expected 1 dependency, got {}", dependencies.len()));
}
let arc_value = dependencies
.pop()
.unwrap()
.downcast::<Option<T>>()
.map_err(|_| {
format!(
"Type mismatch: expected Arc<Option<{}>>",
std::any::type_name::<T>()
)
})?;
Ok((*arc_value).clone())
type Input = (Self,);
type Retv<'input> = &'input Option<T>;
fn extract_from_task_input<'input>(
input: TaskInput<'input, (Self,)>,
) -> Result<Self::Retv<'input>, String> {
Ok(input.next().0)
}
}
@@ -141,77 +122,43 @@ impl<T> ExtractInput for Vec<T>
where
T: 'static + Clone + Send + Sync,
{
async fn extract_from_deps(
mut dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> Result<Self, String> {
if dependencies.len() != 1 {
return Err(format!("Expected 1 dependency, got {}", dependencies.len()));
}
let arc_value = dependencies
.pop()
.unwrap()
.downcast::<Vec<T>>()
.map_err(|_| {
format!(
"Type mismatch: expected Arc<Vec<{}>>",
std::any::type_name::<T>()
)
})?;
Ok((*arc_value).clone())
type Input = (Self,);
type Retv<'input> = &'input Vec<T>;
fn extract_from_task_input<'input>(
input: TaskInput<'input, (Self,)>,
) -> Result<Self::Retv<'input>, String> {
Ok(input.next().0)
}
}
// Implement ExtractInput for HashMap types
impl<K, V> ExtractInput for std::collections::HashMap<K, V>
impl<K, V> ExtractInput for HashMap<K, V>
where
K: 'static + Clone + Send + Sync + Eq + std::hash::Hash,
K: 'static + Clone + Send + Sync + Eq + Hash,
V: 'static + Clone + Send + Sync,
{
async fn extract_from_deps(
mut dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> Result<Self, String> {
if dependencies.len() != 1 {
return Err(format!("Expected 1 dependency, got {}", dependencies.len()));
}
let arc_value = dependencies
.pop()
.unwrap()
.downcast::<std::collections::HashMap<K, V>>()
.map_err(|_| {
format!(
"Type mismatch: expected Arc<HashMap<{}, {}>>",
std::any::type_name::<K>(),
std::any::type_name::<V>()
)
})?;
Ok((*arc_value).clone())
type Input = (Self,);
type Retv<'input> = &'input HashMap<K, V>;
fn extract_from_task_input<'input>(
input: TaskInput<'input, (Self,)>,
) -> Result<Self::Retv<'input>, String> {
Ok(input.next().0)
}
}
// Implement ExtractInput for Arc types - SPECIAL CASE: just clone the Arc, don't unwrap!
// This allows zero-copy data sharing for users who wrap their data in Arc.
impl<T> ExtractInput for Arc<T>
where
T: 'static + Clone + Send + Sync,
{
async fn extract_from_deps(
mut dependencies: Vec<Arc<dyn Any + Send + Sync>>,
) -> Result<Self, String> {
if dependencies.len() != 1 {
return Err(format!("Expected 1 dependency, got {}", dependencies.len()));
}
let arc_arc_value = dependencies
.pop()
.unwrap()
.downcast::<Arc<T>>()
.map_err(|_| {
format!(
"Type mismatch: expected Arc<Arc<{}>>",
std::any::type_name::<T>()
)
})?;
// Return the inner Arc (clone is cheap - just refcount increment)
Ok(Arc::clone(&*arc_arc_value))
type Input = (Self,);
type Retv<'input> = &'input Arc<T>;
fn extract_from_task_input<'input>(
input: TaskInput<'input, (Self,)>,
) -> Result<Self::Retv<'input>, String> {
Ok(input.next().0)
}
}
@@ -228,35 +175,19 @@ where
///
/// Therefore, no runtime validation is needed - if it compiles, the dependency count is correct.
macro_rules! impl_extract_tuple {
($($T:ident : $idx:tt),+) => {
($($T:ident),+) => {
impl<$($T: 'static + Clone + Send + Sync),+> ExtractInput for ($($T,)+) {
async fn extract_from_deps(dependencies: Vec<Arc<dyn Any + Send + Sync>>) -> Result<Self, String> {
// Verify we have the correct number of dependencies
let expected_count = [$($idx),+].len();
if dependencies.len() != expected_count {
return Err(format!("Expected {} dependencies, got {}", expected_count, dependencies.len()));
}
type Input = Self;
type Retv<'input> = ($(&'input $T,)+);
// Convert Vec to iterator so we can consume it
let mut iter = dependencies.into_iter();
fn extract_from_task_input<'input>(input: TaskInput<'input, Self::Input>) -> Result<Self::Retv<'input>, String> {
$(
let ($T, input): (&'input $T, _) = input.next();
)+
// dependencies now contain Arc-wrapped values for efficient fanout
// Dereference Box after downcasting to get the parameter itself
// Use type names as variable names to avoid numeric suffix issues
let arc_results = (
$(
iter.next()
.ok_or_else(|| format!("Missing parameter at index {}", $idx))?
.downcast::<$T>()
.map_err(|_| format!("Type mismatch at index {}: expected Arc<{}>",
$idx, std::any::type_name::<$T>()))?,
)+
);
let _input: TaskInput<'input, ()> = input;
// Clone inner values from Arc (cheap for small types, necessary for ownership)
Ok(($(
(*arc_results.$idx).clone()
,)+))
Ok(($($T,)+))
}
}
};
@@ -265,16 +196,15 @@ macro_rules! impl_extract_tuple {
// Generate ExtractInput implementations for tuples of size 2-8.
// Each invocation creates an impl for a specific tuple size.
// Supporting up to 8 elements covers the vast majority of use cases.
impl_extract_tuple!(A:0, B:1);
impl_extract_tuple!(A:0, B:1, C:2);
impl_extract_tuple!(A:0, B:1, C:2, D:3);
impl_extract_tuple!(A:0, B:1, C:2, D:3, E:4);
impl_extract_tuple!(A:0, B:1, C:2, D:3, E:4, F:5);
impl_extract_tuple!(A:0, B:1, C:2, D:3, E:4, F:5, G:6);
impl_extract_tuple!(A:0, B:1, C:2, D:3, E:4, F:5, G:6, H:7);
#[allow(non_snake_case)]
mod impls {
use super::*;
#[cfg(test)]
mod tests;
#[cfg(test)]
mod coverage_tests;
impl_extract_tuple!(A, B);
impl_extract_tuple!(A, B, C);
impl_extract_tuple!(A, B, C, D);
impl_extract_tuple!(A, B, C, D, E);
impl_extract_tuple!(A, B, C, D, E, F);
impl_extract_tuple!(A, B, C, D, E, F, G);
impl_extract_tuple!(A, B, C, D, E, F, G, H);
}
src/lib.rs
+2 -2
View File
@@ -645,7 +645,7 @@
//! }
//!
//! // Downstream tasks receive &Vec<String> as normal
//! // Behind the scenes: Arc<Vec<String>> is cloned cheaply, then inner Vec extracted
//! // Behind the scenes: Arc<Vec<String>> is cloned cheaply, then a reference to the inner Vec is extracted
//! struct ProcessData;
//! #[task]
//! impl ProcessData {
@@ -791,7 +791,7 @@ mod types;
pub use builder::TaskBuilder;
pub use error::{DagError, DagResult};
pub use runner::DagRunner;
pub use task::Task;
pub use task::{Task, TaskInput};
pub use types::{Pending, TaskHandle};
// Re-export the procedural macro
src/node.rs
+6 -15
View File
@@ -12,8 +12,7 @@ use std::pin::Pin;
use std::sync::Arc;
use crate::error::DagError;
use crate::task::Task;
use crate::types::NodeId;
use crate::task::{Task, TaskInput};
/// Type-erased async execution result (Arc-wrapped for efficient fanout)
type ExecuteFuture =
@@ -52,20 +51,19 @@ pub(crate) trait ExecutableNode: Send {
///
/// The output is also returned from execute_with_deps for central storage.
pub(crate) struct TypedNode<T: Task> {
pub(crate) id: NodeId,
pub(crate) task: T,
}
impl<T: Task + 'static> TypedNode<T> {
pub(crate) fn new(id: NodeId, task: T) -> Self {
Self { id, task }
pub(crate) fn new(task: T) -> Self {
Self { task }
}
}
impl<T: Task + 'static> ExecutableNode for TypedNode<T>
where
T::Input: 'static + Clone,
T::Output: 'static + Clone,
T::Input: 'static,
T::Output: 'static,
{
fn execute_with_deps(
self: Box<Self>,
@@ -75,14 +73,7 @@ where
// 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(dependencies).await {
Ok(output) => output,
Err(_msg) => {
// Failed to extract input - likely type mismatch.
// This should technically be infallible, so it might be worth removing in the future.
return Err(DagError::InvalidDependency { task_id: self.id.0 });
}
};
let output = self.task.run(TaskInput::new(dependencies.iter())).await;
// Arc Wrapping Strategy:
// Wrap output in Arc ONCE for efficient fan-out sharing.
src/runner.rs
+15 -14
View File
@@ -199,8 +199,8 @@ impl DagRunner {
pub fn add_task<Tk>(&self, task: Tk) -> TaskBuilder<'_, Tk, Pending>
where
Tk: Task + Sync + 'static,
Tk::Input: 'static + Clone,
Tk::Output: 'static + Clone,
Tk::Input: 'static,
Tk::Output: 'static,
{
let id = self.alloc_id();
@@ -211,7 +211,7 @@ impl DagRunner {
"adding task to DAG"
);
let node = TypedNode::new(id, task);
let node = TypedNode::new(task);
self.nodes.lock().push(Some(Box::new(node)));
self.edges.lock().insert(id, Vec::new());
self.dependents.lock().insert(id, Vec::new());
@@ -519,6 +519,7 @@ impl DagRunner {
/// # Behavior
///
/// All task outputs are stored after execution and can be retrieved via get().
/// Each task's output can only be taken once, after which it will return `DagError::ResultNotFound`.
///
/// # Errors
///
@@ -526,6 +527,9 @@ impl DagRunner {
/// - The task hasn't been executed yet
/// - The handle is invalid
///
/// Returns `DagError::TaskPanicked` if:
/// - A task executing in the
///
/// # Examples
///
/// ```no_run
@@ -557,26 +561,26 @@ impl DagRunner {
/// assert_eq!(dag.get(task).unwrap(), 42);
/// # };
/// ```
pub fn get<T: 'static + Clone + Send + Sync, H>(&self, handle: H) -> DagResult<T>
pub fn get<T: 'static + Send + Sync, H>(&self, handle: H) -> DagResult<T>
where
H: Into<TaskHandle<T>>,
{
let handle: TaskHandle<T> = handle.into();
let outputs = self.outputs.lock();
let mut outputs = self.outputs.lock();
let arc_output = outputs.get(&handle.id).ok_or(DagError::ResultNotFound {
let arc_output = outputs.remove(&handle.id).ok_or(DagError::ResultNotFound {
task_id: handle.id.0,
})?;
// Downcast Arc<dyn Any> to Arc<T>, then clone the inner value
// The Arc itself is stored for efficient sharing, but get() clones the data
Arc::clone(arc_output)
// Downcast Arc<dyn Any> to Arc<T>
let output = arc_output
.downcast::<T>()
.map(|arc| (*arc).clone())
.map_err(|_| DagError::TypeMismatch {
expected: std::any::type_name::<T>(),
found: "unknown",
})
})?;
Ok(Arc::into_inner(output).unwrap())
}
#[inline]
@@ -673,6 +677,3 @@ impl DagRunner {
#[cfg(test)]
mod tests;
#[cfg(test)]
mod coverage_tests;
src/task.rs
+82 -54
View File
@@ -3,7 +3,9 @@
//! This module defines the core Task trait that all DAG nodes must implement,
//! along with convenience functions for creating tasks from closures.
use std::{future::Future, sync::Arc};
use std::{any::Any, future::Future, marker::PhantomData, slice::Iter, sync::Arc};
use crate::extract::ExtractInput;
/// A unit of async work with typed inputs and outputs.
///
@@ -96,18 +98,82 @@ pub trait Task: Send {
type Input: Send;
type Output: Send + Sync; // Sync required for Arc-wrapping and cross-thread sharing
fn run(self, input: Self::Input) -> impl Future<Output = Self::Output> + Send;
fn run(self, input: TaskInput<Self::Input>) -> impl Future<Output = Self::Output> + Send;
}
/// Internal method: Extract input from list of type-erased dependencies 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,
dependencies: Vec<Arc<dyn std::any::Any + Send + Sync>>,
) -> impl Future<Output = Result<Self::Output, String>> + Send;
pub struct TaskInput<'inputs, Input: Send> {
inputs: Iter<'inputs, Arc<dyn Any + Send + Sync + 'static>>,
phantom: PhantomData<Input>,
}
impl<'inputs, Input: Send> TaskInput<'inputs, Input> {
pub(crate) fn new(inputs: Iter<'inputs, Arc<dyn Any + Send + Sync + 'static>>) -> Self {
Self {
inputs,
phantom: PhantomData,
}
}
}
macro_rules! impl_task_input {
($First:ident $(, $T:ident)*) => {
impl<'inputs, $First: 'static + Send, $($T: 'static + Send),*> TaskInput<'inputs, ($First, $($T,)*)> {
#[must_use]
pub fn next(mut self) -> (&'inputs $First, TaskInput<'inputs, ($($T,)*)>) {
let value = self.inputs.next().and_then(|value| value.downcast_ref()).unwrap();
let next_inputs = TaskInput {
inputs: self.inputs,
phantom: PhantomData,
};
(value, next_inputs)
}
}
};
}
impl_task_input!(A);
impl_task_input!(A, B);
impl_task_input!(A, B, C);
impl_task_input!(A, B, C, D);
impl_task_input!(A, B, C, D, E);
impl_task_input!(A, B, C, D, E, F);
impl_task_input!(A, B, C, D, E, F, G);
impl_task_input!(A, B, C, D, E, F, G, H);
impl TaskInput<'static, ()> {
pub fn empty() -> Self {
Self {
inputs: [].iter(),
phantom: PhantomData,
}
}
}
/// **INTERNAL USE ONLY**: This struct is for testing purposes only and should
/// not be used in production code. Use the `#[task]` macro instead.
#[doc(hidden)]
pub struct TaskFn<I, O, F>
where
for<'input> F: FnMut(I::Retv<'input>) -> O + Send,
I: ExtractInput + Send + Sync + 'static,
O: Send,
{
f: F,
_phantom: std::marker::PhantomData<fn(I) -> O>,
}
impl<I, O, F> Task for TaskFn<I, O, F>
where
for<'input> F: FnMut(I::Retv<'input>) -> O + Send,
I: ExtractInput + Send + Sync + 'static,
O: Send + Sync,
{
type Input = I::Input;
type Output = O;
async fn run(mut self, input: TaskInput<'_, I::Input>) -> O {
(self.f)(I::extract_from_task_input(input).unwrap())
}
}
/// Convenience function to create a task from a closure.
@@ -115,51 +181,13 @@ pub trait Task: Send {
/// **INTERNAL USE ONLY**: This function is for testing purposes only and should
/// not be used in production code. Use the `#[task]` macro instead.
#[doc(hidden)]
pub fn task_fn<I, O, Fut, F>(f: F) -> impl Task<Input = I, Output = O>
pub fn task_fn<I, O, F>(f: F) -> TaskFn<I, O, F>
where
F: FnMut(I) -> Fut + Send + 'static,
Fut: Future<Output = O> + Send + 'static,
I: Send + crate::extract::ExtractInput + 'static,
for<'input> F: FnMut(I::Retv<'input>) -> O + Send,
I: ExtractInput + Send + Sync + 'static,
I::Input: Send + Sync + 'static,
O: Send + Sync + 'static,
{
struct TaskFn<I, O, Fut, F>
where
F: FnMut(I) -> Fut + Send,
Fut: Future<Output = O> + Send,
I: Send,
O: Send,
{
f: F,
_phantom: std::marker::PhantomData<fn(I) -> O>,
}
impl<I, O, Fut, F> Task for TaskFn<I, O, Fut, F>
where
F: FnMut(I) -> Fut + Send,
Fut: Future<Output = O> + Send,
I: Send + crate::extract::ExtractInput,
O: Send + Sync,
{
type Input = I;
type Output = O;
async fn run(mut self, input: I) -> O {
(self.f)(input).await
}
async fn extract_and_run(
self,
dependencies: Vec<Arc<dyn std::any::Any + Send + Sync>>,
) -> Result<Self::Output, String> {
// For task_fn, we use ExtractInput since we don't have macro-generated extraction
let input = I::extract_from_deps(dependencies)
.await
.map_err(|e| format!("Failed to extract input: {}", e))?;
Ok(self.run(input).await)
}
}
TaskFn {
f,
_phantom: std::marker::PhantomData,
src/types.rs
+2 -5
View File
@@ -3,7 +3,7 @@
//! This module defines the fundamental types used throughout the DAG system,
//! including node identifiers, task handles, and type-state markers.
use crate::builder::{IsUnitType, TaskBuilder};
use crate::builder::TaskBuilder;
use crate::task::Task;
/// Opaque node identifier
@@ -63,10 +63,7 @@ pub struct Pending;
// TaskBuilder can be converted to a TaskHandle ONLY if Input is () (unit type)
// This enforces at compile time that tasks with non-unit inputs MUST call .depends_on()
impl<'a, Tk: Task, Deps> From<TaskBuilder<'a, Tk, Deps>> for TaskHandle<Tk::Output>
where
Tk::Input: IsUnitType,
{
impl<'a, Tk: Task<Input = ()>, Deps> From<TaskBuilder<'a, Tk, Deps>> for TaskHandle<Tk::Output> {
fn from(node: TaskBuilder<'a, Tk, Deps>) -> Self {
TaskHandle {
id: node.id,