mpc-macros 0.2.10

/*!
# Task Definition Macro Implementation

The `define_task!` macro is a procedural macro that generates async-mpc task implementations.
It automates the creation of task wrappers, async execution logic, and dependency management.

## Usage

The macro takes a struct definition and a compute function, and generates:
- An internal unresolved struct to hold task dependencies
- A public type alias using the `TaskWrapper` type
- A Task trait implementation with async execution logic
- A `new` constructor for creating new task instances
- A standalone `compute` method for direct calls

## Basic Syntax

```rust
define_task! {
    [pub/pub(crate)/etc] struct TaskName<Generics> {
        field1: Type1,
        field2: Type2,
        // ... more fields
    }

    async fn compute(
        field1: Type1,
        field2: Type2,
        // ... parameters matching struct fields
    ) -> Result<OutputType, AbortError> {
        // Task implementation
        Ok(result)
    }
}
```

**Visibility Modifiers**: The macro respects any visibility modifier placed before the
struct keyword. The same visibility will be applied to both the generated unresolved
struct and the public type alias. If no visibility modifier is provided, the generated
items will be private (module-local).

## Output Type Inference

The macro automatically infers the output type from the compute function's return type.

## Special Field Types

### Task Dependencies
Fields of type `Arc<dyn Task<Output = Arc<SomeType>>>` are treated as task dependencies:
- Automatically spawned for concurrent execution
- Awaited using `tokio::try_join!` for parallel processing
- Results are passed to compute function as `Arc<SomeType>` or `&SomeType` based on the parameter type

### Future Fields
Fields starting with `Next[...]` (preprocessing futures) are treated as futures:
- Awaited alongside task dependencies
- Results passed directly to compute function

### Network Tasks
Tasks with both `network_interface` and `label` fields get special handling:
- Label field is made mutable (`&mut Label`) in compute function calls
- Enables label flag incrementing for network communication patterns
- This means you can compose network-bound tasks without worrying about label flag management

## Examples

### Simple Field Addition
```rust
define_task! {
    pub(crate) struct FieldAddTask<F: FieldExtension> {
        x: Arc<dyn Task<Output = Arc<FieldShare<F>>>>,
        y: Arc<dyn Task<Output = Arc<FieldShare<F>>>>,
    }

    async fn compute(x: &FieldShare<F>, y: &FieldShare<F>) -> Result<FieldShare<F>, AbortError> {
        Ok(x + y)
    }
}
```

Alternatively, if you need ownership of the underlying value, instead of cloning a reference prefer `Arc::unwrap_or_clone`:
```rust
define_task! {
    pub(crate) struct FieldAddTask<F: FieldExtension> {
        x: Arc<dyn Task<Output = Arc<FieldShare<F>>>>,
        y: Arc<dyn Task<Output = Arc<FieldShare<F>>>>,
    }

    async fn compute(x: Arc<FieldShare<F>>, y: &FieldShare<F>) -> Result<FieldShare<F>, AbortError> {
        // Can clone the Arc if needed for sharing, or unwrap for exclusive access
        let x = Arc::unwrap_or_clone(x);
        Ok(x + y)
    }
}
```

**Expands to:**
```rust
// Internal unresolved struct (same visibility as original)
pub(crate) struct UnresolvedFieldAddTask<F: FieldExtension> {
    x: Arc<dyn Task<Output = Arc<FieldShare<F>>>>,
    y: Arc<dyn Task<Output = Arc<FieldShare<F>>>>,
}

// Public type alias (same visibility as original)
pub(crate) type FieldAddTask<F: FieldExtension> =
    TaskWrapper<UnresolvedFieldAddTask<F>, FieldShare<F>>;

// Task trait implementation with async execution
#[async_trait::async_trait]
impl<F: FieldExtension> Task for FieldAddTask<F> {
    type Output = Arc<FieldShare<F>>;

    async fn execute(&self) -> Result<Self::Output, AbortError> {
        // Spawning, awaiting, and compute logic...
    }
}

// Constructor and standalone compute method...
```

### Network Task with Label
```rust
define_task! {
    struct FieldMultiplyTask<F, S, Interface> {
        label: Label,
        network_interface: Arc<Interface>,
        x: Arc<dyn Task<Output = Arc<FieldShare<F>>>>,
        y: Arc<dyn Task<Output = Arc<FieldShare<F>>>>,
        triple: S::NextOutput<Triple<F>>,
    }

    async fn compute(
        label: &mut Label,
        network_interface: Arc<Interface>,
        x: &FieldShare<F>,
        y: &FieldShare<F>,
        triple: Triple<F>,
    ) -> Result<FieldShare<F>, AbortError> {
        // Secure multiplication using Beaver triples
        // Uses label.increment_and_copy() for network communication
    }
}
```

### Preprocessing Source Task
```rust
define_task! {
    struct RandomSingletTask<F: FieldExtension, S: PreprocessingSource<Singlet<F>>> {
        singlet: S::NextSinglet,
    }

    async fn compute(singlet: Singlet<F>) -> Result<FieldShare<F>, AbortError> {
        let Singlet(field_share) = singlet;
        Ok(field_share)
    }
}
```

## Error Handling

The macro will produce compile-time errors for:
- Malformed compute function signatures
- Missing `Result<T, AbortError>` return type
*/

use proc_macro::TokenStream;
use proc_macro2::{Ident, TokenStream as TokenStream2};
use quote::{format_ident, quote};
use syn::{
    parse::{Parse, ParseStream},
    parse_macro_input,
    spanned::Spanned,
    Fields,
    FieldsNamed,
    GenericParam,
    Generics,
    ItemFn,
    ItemStruct,
    Type,
    TypePath,
};

pub fn define_task(input: TokenStream) -> TokenStream {
    let task_def = parse_macro_input!(input as TaskDefinition);
    match expand_task_definition(task_def) {
        Ok(result) => result.into(),
        Err(err) => err.to_compile_error().into(),
    }
}

struct TaskDefinition {
    struct_def: ItemStruct,
    compute_fn: ItemFn,
}

impl Parse for TaskDefinition {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let struct_def: ItemStruct = input.parse()?;
        let compute_fn: ItemFn = input.parse()?;

        Ok(TaskDefinition {
            struct_def,
            compute_fn,
        })
    }
}

fn expand_task_definition(task_def: TaskDefinition) -> syn::Result<TokenStream2> {
    let TaskDefinition {
        struct_def,
        compute_fn,
    } = task_def;

    // Extract output type from compute function's return type
    let output_type = extract_output_type_from_compute_fn(&compute_fn)?;

    let struct_name = &struct_def.ident;
    let generics = &struct_def.generics;
    let visibility = &struct_def.vis;
    let fields = match &struct_def.fields {
        Fields::Named(fields) => fields,
        _ => {
            return Err(syn::Error::new(
                struct_def.span(),
                "Only named fields supported",
            ))
        }
    };

    // Detect if this is a network task by checking for network_interface and label fields
    let network = has_network_fields(fields);

    let unresolved_struct_name = format_ident!("Unresolved{}", struct_name);
    let task_name = struct_name;

    // Use generic TaskWrapper for all tasks
    let wrapper_type = quote! { crate::tasks::TaskWrapper };

    // The output type for Task should be Arc<output> (e.g., Arc<FieldShare<F>>)
    let arc_output_type: Type = syn::parse2(quote! { Arc<#output_type> })?;

    // Generate the code
    let unresolved_struct = generate_unresolved_struct_new(
        &unresolved_struct_name,
        generics,
        fields,
        network,
        visibility,
    );

    let type_alias = generate_type_alias_new(
        task_name,
        &unresolved_struct_name,
        generics,
        &wrapper_type,
        &output_type, // This is the full output type (e.g., FieldShare<F>)
        visibility,
    );

    let task_impl = generate_task_impl_new(
        task_name,
        &unresolved_struct_name,
        generics,
        fields,
        &compute_fn,
        &arc_output_type,
        network,
    )?;

    let constructor = generate_constructor_new(
        task_name,
        &unresolved_struct_name,
        generics,
        fields,
        network,
        visibility,
    );

    let standalone_compute =
        generate_standalone_compute_new(task_name, &compute_fn, generics, visibility);

    Ok(quote! {
        #unresolved_struct
        #type_alias
        #task_impl
        #constructor
        #standalone_compute
    })
}

/// Generates the internal unresolved struct that holds task dependencies before execution
/// This struct contains all the task inputs in their raw form (before awaiting/spawning)
fn generate_unresolved_struct_new(
    name: &Ident,
    generics: &Generics,
    fields: &FieldsNamed,
    _network: bool,
    visibility: &syn::Visibility,
) -> TokenStream2 {
    let (impl_generics, _ty_generics, where_clause) = generics.split_for_impl();

    let mut struct_fields = Vec::new();

    // Add all fields as they are - no special network field handling
    for field in &fields.named {
        let field_name = &field.ident;
        let field_type = &field.ty;
        struct_fields.push(quote! {
            #field_name: #field_type,
        });
    }

    quote! {
        #visibility struct #name #impl_generics #where_clause {
            #(#struct_fields)*
        }
    }
}

/// Generates generics for type aliases that handle defaults properly
fn generate_type_alias_generics(generics: &Generics) -> (TokenStream2, TokenStream2) {
    let mut type_params = Vec::new();
    let mut type_args = Vec::new();

    for param in &generics.params {
        match param {
            GenericParam::Type(type_param) => {
                let ident = &type_param.ident;
                let bounds = &type_param.bounds;
                let default = &type_param.default;

                // For type parameters, include bounds and defaults in the declaration
                if bounds.is_empty() && default.is_none() {
                    type_params.push(quote! { #ident });
                } else if bounds.is_empty() && default.is_some() {
                    type_params.push(quote! { #ident = #default });
                } else if default.is_none() {
                    type_params.push(quote! { #ident: #bounds });
                } else {
                    type_params.push(quote! { #ident: #bounds = #default });
                }

                // For type arguments, just use the identifier
                type_args.push(quote! { #ident });
            }
            GenericParam::Lifetime(lifetime_param) => {
                let lifetime = &lifetime_param.lifetime;
                let bounds = &lifetime_param.bounds;

                if bounds.is_empty() {
                    type_params.push(quote! { #lifetime });
                } else {
                    type_params.push(quote! { #lifetime: #bounds });
                }

                type_args.push(quote! { #lifetime });
            }
            GenericParam::Const(const_param) => {
                let ident = &const_param.ident;
                let ty = &const_param.ty;
                let default = &const_param.default;

                if default.is_none() {
                    type_params.push(quote! { const #ident: #ty });
                } else {
                    type_params.push(quote! { const #ident: #ty = #default });
                }

                type_args.push(quote! { #ident });
            }
        }
    }

    let impl_generics = if type_params.is_empty() {
        quote! {}
    } else {
        quote! { <#(#type_params),*> }
    };

    let ty_generics = if type_args.is_empty() {
        quote! {}
    } else {
        quote! { <#(#type_args),*> }
    };

    (impl_generics, ty_generics)
}

/// Generates the public type alias that combines the unresolved struct with a task wrapper
/// This creates the final task type that users interact with (e.g., FieldAddTask<F>)
fn generate_type_alias_new(
    task_name: &Ident,
    unresolved_name: &Ident,
    generics: &Generics,
    wrapper_type: &TokenStream2,
    output_type: &Type,
    visibility: &syn::Visibility,
) -> TokenStream2 {
    let (impl_generics, ty_generics) = generate_type_alias_generics(generics);

    quote! {
        #[allow(type_alias_bounds)]
        #visibility type #task_name #impl_generics = #wrapper_type<#unresolved_name #ty_generics, #output_type>;
    }
}

/// Generates the Task trait implementation with async execution logic
/// Handles spawning task dependencies, awaiting futures, and calling the compute function
/// Includes special handling for mutable label references and concurrent task execution
fn generate_task_impl_new(
    task_name: &Ident,
    unresolved_name: &Ident,
    generics: &Generics,
    fields: &FieldsNamed,
    compute_fn: &ItemFn,
    output_type: &Type,
    _network: bool,
) -> syn::Result<TokenStream2> {
    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

    // Build a map from field names to compute parameter types
    // This allows us to check if a field should be passed as Arc<T> or &T
    let mut field_param_types = std::collections::HashMap::new();
    for input in &compute_fn.sig.inputs {
        if let syn::FnArg::Typed(pat_type) = input {
            if let syn::Pat::Ident(pat_ident) = &*pat_type.pat {
                let param_name = pat_ident.ident.clone();
                let param_type = (*pat_type.ty).clone();
                field_param_types.insert(param_name.to_string(), param_type);
            }
        }
    }

    // Generate field destructuring
    let mut field_names = Vec::new();

    for field in &fields.named {
        let field_name = &field.ident;
        // Make label field mutable in the destructuring pattern
        if *field_name.as_ref().unwrap() == "label" {
            field_names.push(quote! { mut #field_name });
        } else {
            field_names.push(quote! { #field_name });
        }
    }

    // Generate task spawning for Arc<dyn Task<...>> fields
    let mut task_spawns = Vec::new();
    let mut task_awaitable_names = Vec::new();
    let mut vec_task_names = Vec::new();
    let mut vec_future_names = Vec::new();
    let mut single_future_names = Vec::new();
    let mut compute_args = Vec::new();

    for field in &fields.named {
        let field_name = field.ident.as_ref().unwrap();
        let field_type = &field.ty;

        // Check if this is a Vec<Arc<dyn Task<...>>> field
        if is_vec_task_field(field_type) {
            vec_task_names.push(quote! { #field_name });
            // Check if the compute function expects a slice reference (&[T]) instead of Vec<T>
            if let Some(param_type) = field_param_types.get(&field_name.to_string()) {
                if is_slice_reference_param(param_type) {
                    // Pass as slice reference
                    compute_args.push(quote! { &#field_name[..] });
                } else {
                    // Pass as Vec<T> or &Vec<T> depending on parameter type
                    compute_args.push(quote! { #field_name });
                }
            } else {
                // Fallback: pass as Vec<T> if we can't determine the type
                compute_args.push(quote! { #field_name });
            }
        } else if is_task_field_new(field_type) {
            // Check if this is an Arc<dyn Task<...>> field
            task_spawns.push(quote! {
                let #field_name = tokio::task::spawn(async move { #field_name.execute().await });
            });
            task_awaitable_names.push(quote! { #field_name });

            // Check if the compute function expects Arc<T>, &T, or &[T]
            if let Some(param_type) = field_param_types.get(&field_name.to_string()) {
                if is_slice_reference_param(param_type) {
                    // Pass as slice reference for Arc<Vec<T>> -> &[T]
                    compute_args.push(quote! { &std::sync::Arc::as_ref(&#field_name)[..] });
                } else if is_reference_param(param_type) {
                    // Pass as &T using Arc::as_ref
                    compute_args.push(quote! { std::sync::Arc::as_ref(&#field_name) });
                } else {
                    // Pass as Arc<T> directly
                    compute_args.push(quote! { #field_name });
                }
            } else {
                // Fallback: pass as Arc<T> if we can't determine the type
                compute_args.push(quote! { #field_name });
            }
        } else if is_vec_future_field(field_type) {
            // For vectors of futures that need to be awaited (like Vec<NextTriple>)
            vec_future_names.push(quote! { #field_name });
            compute_args.push(quote! { #field_name });
        } else if is_future_field_new(field_type) {
            // For futures that need to be awaited (like Singlet, Triple)
            single_future_names.push(quote! { #field_name });
            compute_args.push(quote! { #field_name });
        } else if is_phantom_data_field(field_type) {
            // Skip PhantomData fields - they don't get passed to compute method
            continue;
        } else {
            // For regular fields that don't need awaiting
            // Special case: pass label as &mut
            if *field_name == "label" {
                compute_args.push(quote! { &mut #field_name });
            } else {
                compute_args.push(quote! { #field_name });
            }
        }
    }

    // Generate the execute method
    let compute_call = quote! {
        let result = Self::compute(#(#compute_args),*).await?;
    };

    // Build awaiting logic based on different types of awaitables
    let mut awaiting_stmts = Vec::new();

    // Handle vector tasks (spawn all and then await all)
    for vec_task_var in &vec_task_names {
        awaiting_stmts.push(quote! {
            let #vec_task_var: Vec<_> = #vec_task_var.into_iter()
                .map(|task| tokio::task::spawn(async move { task.execute().await }))
                .collect();
            let #vec_task_var: Result<Vec<_>, _> = futures::future::try_join_all(
                #vec_task_var.into_iter().map(|handle| async move {
                    handle.await.map_err(Into::<core_utils::errors::AbortError>::into)
                })
            ).await;
            let #vec_task_var: Vec<_> = #vec_task_var?
                .into_iter()
                .collect::<Result<Vec<_>, _>>()?
                .into_iter()
                .map(std::sync::Arc::unwrap_or_clone)
                .collect();
        });
    }

    // Handle tasks (which are spawned)
    if !task_awaitable_names.is_empty() {
        if task_awaitable_names.len() == 1 {
            let task_var = &task_awaitable_names[0];
            awaiting_stmts.push(quote! {
                let #task_var = #task_var.await??;
            });
        } else {
            awaiting_stmts.push(quote! {
                let (#(#task_awaitable_names),*) = tokio::try_join!(
                    #(
                        futures::TryFutureExt::map_err(#task_awaitable_names, Into::<core_utils::errors::AbortError>::into)
                    ),*
                )?;
                let (#(#task_awaitable_names),*) = (#(#task_awaitable_names?),*);
            });
        }
    }

    // Handle vector futures
    for vec_future_var in &vec_future_names {
        awaiting_stmts.push(quote! {
            let #vec_future_var = futures::future::try_join_all(#vec_future_var).await?;
        });
    }

    // Handle single futures
    for single_future_var in &single_future_names {
        awaiting_stmts.push(quote! {
            let #single_future_var = #single_future_var.await?;
        });
    }

    let awaiting_and_unwrapping = quote! {
        #(#awaiting_stmts)*
    };

    Ok(quote! {
        #[async_trait::async_trait]
        impl #impl_generics crate::tasks::Task for #task_name #ty_generics #where_clause {
            type Output = #output_type;

            async fn execute(&self) -> Result<Self::Output, core_utils::errors::AbortError> {
                let mut state = self.state.lock().await;
                if let Some(unresolved_state) = state.take_unresolved()? {
                    let #unresolved_name { #(#field_names),* } = unresolved_state;

                    #(#task_spawns)*

                    #awaiting_and_unwrapping

                    #compute_call
                    state.resolve(std::sync::Arc::new(result));
                }
                Ok(state.clone_output()?)
            }
        }
    })
}

/// Generates the task constructor that creates a new task instance
/// Returns an Arc<Task> with the unresolved struct wrapped in a mutex for thread safety
fn generate_constructor_new(
    task_name: &Ident,
    unresolved_name: &Ident,
    generics: &Generics,
    fields: &FieldsNamed,
    _network: bool,
    _visibility: &syn::Visibility,
) -> TokenStream2 {
    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();

    // Generate constructor parameters
    let mut constructor_params = Vec::new();
    let mut field_assignments = Vec::new();

    for field in &fields.named {
        let field_name = &field.ident;
        let field_type = &field.ty;

        if is_phantom_data_field(field_type) {
            // For PhantomData fields, provide Default::default() and don't include in constructor
            // parameters
            field_assignments.push(quote! { #field_name: Default::default() });
        } else {
            // For regular fields, include in constructor parameters
            constructor_params.push(quote! { #field_name: #field_type });
            field_assignments.push(quote! { #field_name });
        }
    }

    quote! {
        impl #impl_generics #task_name #ty_generics #where_clause {
            pub fn new(#(#constructor_params),*) -> std::sync::Arc<Self> {
                let unresolved = #unresolved_name {
                    #(#field_assignments),*
                };
                std::sync::Arc::new(Self {
                    state: tokio::sync::Mutex::new(crate::tasks::InternalTaskState::new(unresolved)),
                })
            }
        }
    }
}

/// Generates a standalone compute method that can be called directly without the task wrapper
/// Useful for testing and direct computation without the async task infrastructure
fn generate_standalone_compute_new(
    task_name: &Ident,
    compute_fn: &ItemFn,
    generics: &Generics,
    _visibility: &syn::Visibility,
) -> TokenStream2 {
    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
    let sig = &compute_fn.sig;
    let block = &compute_fn.block;

    quote! {
        impl #impl_generics #task_name #ty_generics #where_clause {
            pub #sig #block
        }
    }
}

/// Checks if a field type represents a task dependency (Arc<dyn Task<Output = ...>>)
/// These fields need special handling as they must be spawned and awaited concurrently
fn is_task_field_new(ty: &Type) -> bool {
    // Check if the type is Arc<dyn Task<Output = ...>>
    if let Type::Path(TypePath { path, .. }) = ty {
        if let Some(segment) = path.segments.last() {
            if segment.ident == "Arc" {
                // Look for "Task" in the type arguments
                if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
                    for arg in &args.args {
                        if let syn::GenericArgument::Type(Type::TraitObject(trait_obj)) = arg {
                            for bound in &trait_obj.bounds {
                                if let syn::TypeParamBound::Trait(trait_bound) = bound {
                                    if trait_bound.path.segments.last().unwrap().ident == "Task" {
                                        return true;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    false
}

/// Checks if a field type represents a future that needs to be awaited
/// These are typically preprocessing sources that yield futures (e.g., NextSinglet<F>>)
/// Also handles vectors of futures (e.g., Vec<NextTriple>)
fn is_future_field_new(ty: &Type) -> bool {
    // Check if the type contains "Next" which indicates a future
    if let Type::Path(TypePath { path, .. }) = ty {
        for segment in &path.segments {
            if segment.ident.to_string().contains("Next") {
                return true;
            }
            // Check if this is a Vec<...> and the inner type contains "Next"
            if segment.ident == "Vec" {
                if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
                    for arg in &args.args {
                        if let syn::GenericArgument::Type(inner_type) = arg {
                            if is_future_field_new(inner_type) {
                                return true;
                            }
                        }
                    }
                }
            }
        }
    }
    false
}

/// Checks if a field type is a vector of futures (e.g., Vec<NextTriple>)
fn is_vec_future_field(ty: &Type) -> bool {
    if let Type::Path(TypePath { path, .. }) = ty {
        if let Some(segment) = path.segments.last() {
            if segment.ident == "Vec" {
                if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
                    for arg in &args.args {
                        if let syn::GenericArgument::Type(inner_type) = arg {
                            if is_future_field_new(inner_type) {
                                return true;
                            }
                        }
                    }
                }
            }
        }
    }
    false
}

/// Checks if a field type is a vector of tasks (e.g., Vec<Arc<dyn Task<...>>>)
fn is_vec_task_field(ty: &Type) -> bool {
    if let Type::Path(TypePath { path, .. }) = ty {
        if let Some(segment) = path.segments.last() {
            if segment.ident == "Vec" {
                if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
                    for arg in &args.args {
                        if let syn::GenericArgument::Type(inner_type) = arg {
                            if is_task_field_new(inner_type) {
                                return true;
                            }
                        }
                    }
                }
            }
        }
    }
    false
}

/// Checks if a field type is PhantomData and should be ignored
/// PhantomData fields are used for type-level programming but don't contain actual data
fn is_phantom_data_field(ty: &Type) -> bool {
    if let Type::Path(TypePath { path, .. }) = ty {
        if let Some(segment) = path.segments.last() {
            return segment.ident == "PhantomData";
        }
    }
    false
}

/// Determines if a task struct has network-related fields (network_interface and label)
/// This affects code generation for network communication patterns
fn has_network_fields(fields: &FieldsNamed) -> bool {
    let mut has_network_interface = false;
    let mut has_label = false;

    for field in &fields.named {
        if let Some(field_name) = &field.ident {
            match field_name.to_string().as_str() {
                "network_interface" => has_network_interface = true,
                "label" => has_label = true,
                _ => {}
            }
        }
    }

    has_network_interface && has_label
}

/// Checks if a compute function parameter type is a reference (&T)
fn is_reference_param(ty: &Type) -> bool {
    matches!(ty, Type::Reference(_))
}

/// Checks if a compute function parameter type is a slice reference (&[T])
fn is_slice_reference_param(ty: &Type) -> bool {
    if let Type::Reference(type_ref) = ty {
        if let Type::Slice(_) = &*type_ref.elem {
            return true;
        }
    }
    false
}

/// Extracts the output type from a compute function's return type
/// Parses Result<T, AbortError> -> T to automatically infer the task's output type
/// This eliminates the need for explicit output field declarations in the macro
fn extract_output_type_from_compute_fn(compute_fn: &ItemFn) -> syn::Result<Type> {
    // Extract the return type from Result<T, AbortError> -> T
    if let syn::ReturnType::Type(_, return_type) = &compute_fn.sig.output {
        if let syn::Type::Path(type_path) = return_type.as_ref() {
            // Look for Result<T, AbortError> pattern
            if let Some(segment) = type_path.path.segments.last() {
                if segment.ident == "Result" {
                    if let syn::PathArguments::AngleBracketed(args) = &segment.arguments {
                        if let Some(syn::GenericArgument::Type(output_type)) = args.args.first() {
                            return Ok(output_type.clone());
                        }
                    }
                }
            }
        }
    }

    Err(syn::Error::new(
        compute_fn.sig.span(),
        "Could not extract output type from compute function return type",
    ))
}