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

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,