state-macro 0.1.1

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::{
    Expr, ExprCall, ExprPath, Result, Token,
    parse::{Parse, ParseStream},
    parse_macro_input,
    punctuated::Punctuated,
    visit_mut::{self, VisitMut},
};

/// Parses the content of the `with_state!` macro.
struct WithStateInput {
    state_expr: Expr,
    content: TokenStream2,
}

impl Parse for WithStateInput {
    fn parse(input: ParseStream) -> Result<Self> {
        // Parse the state expression
        let state_expr = input.parse::<Expr>()?;
        // Expect semicolon after state expression
        input.parse::<Token![;]>()?;

        // The rest is the content that we'll transform
        let content = input.parse()?;

        Ok(WithStateInput {
            state_expr,
            content,
        })
    }
}

/// Visitor that transforms function calls with paths starting with `::`
struct StatefulCallVisitor<'a> {
    state_expr: &'a Expr,
}

impl VisitMut for StatefulCallVisitor<'_> {
    fn visit_expr_call_mut(&mut self, call: &mut ExprCall) {
        // First visit any nested expressions in the function call
        visit_mut::visit_expr_call_mut(self, call);

        // Check if the function path starts with `::`
        if let Expr::Path(ExprPath { path, .. }) = &mut *call.func {
            if path.leading_colon.is_some() {
                // This is a stateful function call, so we need to add the state as first arg

                // Remove the leading :: from the path
                let mut new_path = path.clone();
                new_path.leading_colon = None;

                // Update the function path without the leading ::
                call.func = Box::new(Expr::Path(ExprPath {
                    attrs: vec![],
                    qself: None,
                    path: new_path,
                }));

                // Prepare new arguments with state_expr as first arg
                let mut new_args = Punctuated::new();
                new_args.push(self.state_expr.clone());

                // Add existing args
                for arg in call.args.iter() {
                    new_args.push(arg.clone());
                }

                // Replace the arguments
                call.args = new_args;
            }
        }
    }
}

/// The with_state! macro implementation
pub fn with_state(input: TokenStream) -> TokenStream {
    let WithStateInput {
        state_expr,
        content,
    } = parse_macro_input!(input as WithStateInput);

    // Parse the content as a block of statements
    let mut parsed_content = match syn::parse2::<syn::Block>(quote! { { #content } }) {
        Ok(block) => block.stmts,
        Err(_) => {
            // If it can't be parsed as a block, return a compile error
            return TokenStream::from(quote! {
                compile_error!("Failed to parse macro content")
            });
        }
    };

    // Create a visitor to transform stateful function calls
    let mut visitor = StatefulCallVisitor {
        state_expr: &state_expr,
    };

    // Process each statement
    for stmt in &mut parsed_content {
        visitor.visit_stmt_mut(stmt);
    }

    // Return the transformed statements directly (without block braces)
    TokenStream::from(quote! {
        #(#parsed_content)*
    })
}