type_flow_proc_macros/

lib.rs

use proc_macro::TokenStream;
use quote::{quote, format_ident};
use syn::{parse_macro_input, Ident, Type, parse::Parse, parse::ParseStream, Token, /*GenericParam,*/ ItemStruct, FieldsNamed, Fields};

struct ProcessorPipeline {
    pipeline_name: Ident,
    data_type: Type,
    processors: Vec<Type>,
}

impl Parse for ProcessorPipeline {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let pipeline_name: Ident = input.parse()?;
        input.parse::<Token![,]>()?;
        let data_type: Type = input.parse()?;
        input.parse::<Token![,]>()?;
        let mut processors = Vec::new();
        // Parse the remaining types as processor types
        while !input.is_empty() {
            let processor_type: Type = input.parse()?;
            processors.push(processor_type);
            if input.is_empty() {
                break;
            }
            input.parse::<Token![,]>()?;
        }
        Ok(ProcessorPipeline {
            pipeline_name,
            data_type,
            processors,
        })
    }
}

#[proc_macro]
pub fn stateful_processor_pipeline_with_index(input: TokenStream) -> TokenStream {
    let ProcessorPipeline { pipeline_name, data_type, processors } = parse_macro_input!(input as ProcessorPipeline);
    let struct_fields = processors.iter().enumerate().map(|(idx, processor_type)| {
        // Get a string representation of the type for the field name
        let type_str = get_type_name(processor_type);
        let field_name = format_ident!("processor_{}_{}", type_str, idx);
        quote! { #field_name: #processor_type }
    });
    let constructor_params = processors.iter().enumerate().map(|(idx, processor_type)| {
        let type_str = get_type_name(processor_type);
        let param_name = format_ident!("processor_{}_{}", type_str, idx);
        quote! { #param_name: #processor_type }
    });
    let field_initializers = processors.iter().enumerate().map(|(idx, processor_type)| {
        let type_str = get_type_name(processor_type);
        let field_name = format_ident!("processor_{}_{}", type_str, idx);
        quote! { #field_name }
    });
    let process_implementation = processors.iter().enumerate().map(|(idx, processor_type)| {
        let type_str = get_type_name(processor_type);
        let field_name = format_ident!("processor_{}_{}", type_str, idx);
        quote! { data = self.#field_name.process(data); }
    });
    let expanded = quote! {
        pub struct #pipeline_name {
            #(#struct_fields,)*
        }
        impl #pipeline_name {
            pub fn new(#(#constructor_params,)*) -> Self {
                Self {
                    #(#field_initializers,)*
                }
            }
        }
        impl crate::StatefulProcessor<#data_type> for #pipeline_name {
            fn process(&mut self, mut data: #data_type) -> #data_type {
                #(#process_implementation)*
                data
            }
        }
    };
    TokenStream::from(expanded)
}

struct ProcessorInplacePipeline {
    pipeline_name: Ident,
    data_type: Type,
    error_type: Type,
    processors: Vec<Type>,
}

impl Parse for ProcessorInplacePipeline {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let pipeline_name: Ident = input.parse()?;
        input.parse::<Token![,]>()?;
        let data_type: Type = input.parse()?;
        input.parse::<Token![,]>()?;
        let error_type: Type = input.parse()?;
        input.parse::<Token![,]>()?;
        let mut processors = Vec::new();
        // Parse the remaining types as processor types
        while !input.is_empty() {
            let processor_type: Type = input.parse()?;
            processors.push(processor_type);

            if input.is_empty() {
                break;
            }
            input.parse::<Token![,]>()?;
        }
        Ok(ProcessorInplacePipeline{
            pipeline_name,
            data_type,
            error_type,
            processors,
        })
    }
}


#[proc_macro_attribute]
pub fn implement_processor_swapping(_attr: proc_macro::TokenStream, item: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input_struct_item = parse_macro_input!(item as ItemStruct);
    let struct_name = &input_struct_item.ident;
    let generics = &input_struct_item.generics;
    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
    let type_param_idents: Vec<&Ident> = generics.type_params().map(|tp| &tp.ident).collect();
    let fields = match &input_struct_item.fields {
        Fields::Named(FieldsNamed { named, .. }) => named,
        _ => panic!("#[implement_processor_swapping] only works on structs with named fields."),
    };
    let processor_field_idents: Vec<&Ident> = fields.iter()
        .filter_map(|f| f.ident.as_ref())
        .filter(|&id| id.to_string() != "_marker")
        .collect();
    if type_param_idents.len() != processor_field_idents.len() {
        panic!(
            "#[implement_processor_swapping] detected a mismatch between the number of generic type parameters ({}) and processor fields ({}). Expected them to be equal.",
            type_param_idents.len(),
            processor_field_idents.len()
        );
    }
    // --- Trait Impls ---
    let mut trait_impls = Vec::new();
    if !type_param_idents.is_empty() {
        // Reverse
        let mut rev_types = type_param_idents.clone();
        rev_types.reverse();
        let mut rev_source_fields = processor_field_idents.clone();
        rev_source_fields.reverse();
        let rev_assignments = processor_field_idents.iter().zip(rev_source_fields.iter()).map(|(dest, src)| {
            quote!{ #dest: self.#src }
        });
        trait_impls.push(quote!{
            impl #impl_generics type_flow_traits::Reverse for #struct_name #ty_generics #where_clause {
                type Output = #struct_name<#(#rev_types),*>;
                fn reverse(self) -> Self::Output {
                    Self::Output {
                        _marker: std::marker::PhantomData,
                        #(#rev_assignments),*
                    }
                }
            }
        });
        // ShiftLeft
        let mut sl_types = type_param_idents.clone();
        sl_types.rotate_left(1);
        let mut sl_source_fields = processor_field_idents.clone();
        sl_source_fields.rotate_left(1);
        let sl_assignments = processor_field_idents.iter().zip(sl_source_fields.iter()).map(|(dest, src)| {
            quote!{ #dest: self.#src }
        });
        trait_impls.push(quote!{
            impl #impl_generics type_flow_traits::ShiftLeft for #struct_name #ty_generics #where_clause {
                type ShiftedLeft = #struct_name<#(#sl_types),*>;
                fn shift_left(self) -> Self::ShiftedLeft {
                    Self::ShiftedLeft {
                        _marker: std::marker::PhantomData,
                        #(#sl_assignments),*
                    }
                }
            }
        });
        // ShiftRight
        let mut sr_types = type_param_idents.clone();
        sr_types.rotate_right(1);
        let mut sr_source_fields = processor_field_idents.clone();
        sr_source_fields.rotate_right(1);
        let sr_assignments = processor_field_idents.iter().zip(sr_source_fields.iter()).map(|(dest, src)| {
            quote!{ #dest: self.#src }
        });
        trait_impls.push(quote!{
            impl #impl_generics type_flow_traits::ShiftRight for #struct_name #ty_generics #where_clause {
                type ShiftedRight = #struct_name<#(#sr_types),*>;
                fn shift_right(self) -> Self::ShiftedRight {
                    Self::ShiftedRight {
                        _marker: std::marker::PhantomData,
                        #(#sr_assignments),*
                    }
                }
            }
        });
        // SwapStartEnd
        let mut sse_types = type_param_idents.clone();
        let sse_length = sse_types.len();
        if sse_length > 1 { sse_types.swap(0, sse_length - 1); }
        let mut sse_source_fields = processor_field_idents.clone();
        let sse_source_length = sse_source_fields.len();
        if sse_source_length > 1 { sse_source_fields.swap(0, sse_source_length - 1); }
        let sse_assignments = processor_field_idents.iter().zip(sse_source_fields.iter()).map(|(dest, src)| {
            quote!{ #dest: self.#src }
        });
        trait_impls.push(quote!{
            impl #impl_generics type_flow_traits::SwapStartEnd for #struct_name #ty_generics #where_clause {
                type Output = #struct_name<#(#sse_types),*>;
                fn swap(self) -> Self::Output {
                    Self::Output {
                        _marker: std::marker::PhantomData,
                        #(#sse_assignments),*
                    }
                }
            }
        });
    }
    let num_processors = type_param_idents.len();
    if num_processors < 2 {
        let trait_def_placeholder = quote! {
            pub trait SwapArbitraryProcessors<const I: usize, const J: usize> where Self: Sized {
                type SwappedOutput;
                fn swap_processors(self) -> Self::SwappedOutput;
            }
        };
        return quote! {
            #input_struct_item
            #trait_def_placeholder
            #(#trait_impls)*
        }.into();
    }
    let mut generated_impls = Vec::new();
    let swap_trait_name = format_ident!("SwapArbitraryProcessors");
    for i in 0..num_processors {
        for j in 0..num_processors {
            if i == j { continue; }
            let mut swapped_generic_params_for_type = type_param_idents.clone();
            swapped_generic_params_for_type.swap(i, j);
            let field_i_name = processor_field_idents[i];
            let field_j_name = processor_field_idents[j];
            let mut current_field_initializers = Vec::new();
            for k in 0..num_processors {
                let current_field_name = processor_field_idents[k];
                if k == i {
                    current_field_initializers.push(quote! { #current_field_name: self.#field_j_name });
                } else if k == j {
                    current_field_initializers.push(quote! { #current_field_name: self.#field_i_name });
                } else {
                    current_field_initializers.push(quote! { #current_field_name: self.#current_field_name });
                }
            }
            generated_impls.push(quote! {
                impl #impl_generics #swap_trait_name<#i, #j> for #struct_name #ty_generics #where_clause {
                    type SwappedOutput = #struct_name<#(#swapped_generic_params_for_type),*>;
                    fn swap_processors(self) -> Self::SwappedOutput {
                        #struct_name {
                            _marker: std::marker::PhantomData,
                            #( #current_field_initializers ),*
                        }
                    }
                }
            });
        }
    }

    let trait_def = quote! {
        pub trait #swap_trait_name<const I: usize, const J: usize>
        where
            Self: Sized,
        {
            type SwappedOutput;
            fn swap_processors(self) -> Self::SwappedOutput;
        }
    };

    let expanded = quote! {
        #input_struct_item
        #trait_def
        #(#generated_impls)*
        #(#trait_impls)*
    };
    proc_macro::TokenStream::from(expanded)
}

#[proc_macro]
pub fn inplace_stateful_processor_pipeline_with_index(input: TokenStream) -> TokenStream { 
    let ProcessorInplacePipeline { pipeline_name, data_type, error_type, processors } = parse_macro_input!(input as ProcessorInplacePipeline);
    let struct_fields = processors.iter().enumerate().map(|(idx, processor_type)| {
        // Get a string representation of the type for the field name
        let type_str = get_type_name(processor_type);
        let field_name = format_ident!("processor_{}_{}", type_str, idx);
        quote! { #field_name: #processor_type }
    });
    let constructor_params = processors.iter().enumerate().map(|(idx, processor_type)| {
        let type_str = get_type_name(processor_type);
        let param_name = format_ident!("processor_{}_{}", type_str, idx);
        quote! { #param_name: #processor_type }
    });
    let field_initializers = processors.iter().enumerate().map(|(idx, processor_type)| {
        let type_str = get_type_name(processor_type);
        let field_name = format_ident!("processor_{}_{}", type_str, idx);
        quote! { #field_name }
    });
    let process_implementation = processors.iter().enumerate().map(|(idx, processor_type)| {
        let type_str = get_type_name(processor_type);
        let field_name = format_ident!("processor_{}_{}", type_str, idx);
        quote! { self.#field_name.process(data)?; }
    });
    let expanded = quote! {
        pub struct #pipeline_name {
            #(#struct_fields,)*
        }
        impl #pipeline_name {
            pub fn new(#(#constructor_params,)*) -> Self {
                Self {
                    #(#field_initializers,)*
                }
            }
        }
        impl crate::InPlaceStatefulProcessor<#data_type, #error_type> for #pipeline_name {
            fn process(&mut self, data: &mut #data_type) -> Result<(), #error_type> {
                #(#process_implementation)*
                Ok(())
            }
        }
    };
    TokenStream::from(expanded)
}

// Helper function to extract a usable identifier from a Type
fn get_type_name(ty: &syn::Type) -> String {
    match ty {
        syn::Type::Path(type_path) if !type_path.path.segments.is_empty() => {
            // Get the last segment of the path (e.g., for std::string::String, get "String")
            let segment = type_path.path.segments.last().unwrap();
            let name = segment.ident.to_string();

            // Filter out any characters that aren't valid in an identifier
            name.chars()
                .map(|c| if c.is_alphanumeric() { c } else { '_' })
                .collect()
        },
        // Handle other types (arrays, references, etc.) by using a generic name
        _ => "unknown_type".to_string(),
    }
}