fieldwork-macro 0.1.0

//! Code generation for the with_! macro.

use proc_macro2::{Span, TokenStream};
use quote::{quote, ToTokens};
use std::collections::HashMap;
use syn::{
    visit_mut::VisitMut, Expr, ExprBlock, ExprPath, Ident, Member, Path, PathSegment,
};

use crate::parse::{FieldMode, MethodDef, WithInput};

/// Collected field information for code generation.
struct FieldInfo {
    mode: FieldMode,
    span: Span,
}

/// Transforms the code block by replacing method calls and field accesses.
struct BodyTransformer<'a> {
    receiver: &'a Expr,
    methods: HashMap<String, &'a MethodDef>,
    field_modes: &'a HashMap<String, FieldInfo>,
}

impl<'a> VisitMut for BodyTransformer<'a> {
    fn visit_expr_mut(&mut self, expr: &mut Expr) {
        match expr {
            // Transform: receiver.method(args) -> inlined method body
            Expr::MethodCall(call) => {
                if self.is_receiver(&call.receiver) {
                    let method_name = call.method.to_string();
                    if let Some(method_def) = self.methods.get(&method_name) {
                        *expr = self.inline_method(call, method_def);
                        return;
                    }
                }
                syn::visit_mut::visit_expr_mut(self, expr);
            }

            // Transform: receiver.field -> field
            Expr::Field(field_expr) => {
                if self.is_receiver(&field_expr.base) {
                    if let Member::Named(field_name) = &field_expr.member {
                        let name_str = field_name.to_string();
                        if self.field_modes.contains_key(&name_str) {
                            *expr = self.make_ident_expr(field_name);
                            return;
                        }
                    }
                }
                syn::visit_mut::visit_expr_mut(self, expr);
            }

            _ => {
                syn::visit_mut::visit_expr_mut(self, expr);
            }
        }
    }
}

impl<'a> BodyTransformer<'a> {
    /// Check if an expression matches the receiver.
    fn is_receiver(&self, expr: &Expr) -> bool {
        expr.to_token_stream().to_string() == self.receiver.to_token_stream().to_string()
    }

    /// Create an identifier expression from an Ident.
    fn make_ident_expr(&self, ident: &Ident) -> Expr {
        Expr::Path(ExprPath {
            attrs: vec![],
            qself: None,
            path: Path {
                leading_colon: None,
                segments: std::iter::once(PathSegment {
                    ident: ident.clone(),
                    arguments: syn::PathArguments::None,
                })
                .collect(),
            },
        })
    }

    /// Inline a method call by substituting its body.
    fn inline_method(&self, call: &syn::ExprMethodCall, method: &MethodDef) -> Expr {
        let mut body = method.body.clone();

        // Build replacement map for extra parameters
        let mut replacer = ParamReplacer {
            replacements: HashMap::new(),
        };

        for (param, arg) in method.extra_params.iter().zip(call.args.iter()) {
            replacer
                .replacements
                .insert(param.name.to_string(), arg.clone());
        }

        // Replace extra parameters in method body
        syn::visit_mut::visit_block_mut(&mut replacer, &mut body);

        // Return as block expression
        Expr::Block(ExprBlock {
            attrs: vec![],
            label: None,
            block: body,
        })
    }
}

/// Replaces parameter identifiers with their corresponding expressions.
struct ParamReplacer {
    replacements: HashMap<String, Expr>,
}

impl VisitMut for ParamReplacer {
    fn visit_expr_mut(&mut self, expr: &mut Expr) {
        if let Expr::Path(path) = expr {
            if path.qself.is_none() && path.path.segments.len() == 1 {
                let name = path.path.segments[0].ident.to_string();
                if let Some(replacement) = self.replacements.get(&name) {
                    *expr = replacement.clone();
                    return;
                }
            }
        }
        syn::visit_mut::visit_expr_mut(self, expr);
    }
}

/// Merge field mode, preferring Mut over Ref.
fn merge_mode(existing: FieldMode, new: FieldMode) -> FieldMode {
    match (existing, new) {
        (FieldMode::Mut, _) | (_, FieldMode::Mut) => FieldMode::Mut,
        (FieldMode::Copy, _) | (_, FieldMode::Copy) => FieldMode::Copy,
        _ => FieldMode::Ref,
    }
}

/// Main expansion function.
pub fn expand(input: WithInput) -> syn::Result<TokenStream> {
    let WithInput {
        receiver,
        methods,
        fields,
        mut body,
    } = input;

    // Collect all fields and their access modes
    let mut field_modes: HashMap<String, FieldInfo> = HashMap::new();

    // Collect fields from method definitions
    for method in &methods {
        for param in &method.field_params {
            let name = param.name.to_string();
            field_modes
                .entry(name)
                .and_modify(|info| {
                    info.mode = merge_mode(info.mode, param.mode);
                })
                .or_insert(FieldInfo {
                    mode: param.mode,
                    span: param.name.span(),
                });
        }
    }

    // Collect fields from fields block
    for field in &fields {
        let name = field.name.to_string();
        field_modes
            .entry(name)
            .and_modify(|info| {
                info.mode = merge_mode(info.mode, field.mode);
            })
            .or_insert(FieldInfo {
                mode: field.mode,
                span: field.name.span(),
            });
    }

    // Build method lookup map
    let methods_map: HashMap<String, &MethodDef> =
        methods.iter().map(|m| (m.name.to_string(), m)).collect();

    // Transform the code block
    let mut transformer = BodyTransformer {
        receiver: &receiver,
        methods: methods_map,
        field_modes: &field_modes,
    };
    syn::visit_mut::visit_block_mut(&mut transformer, &mut body);

    // Generate field bindings
    let bindings: Vec<TokenStream> = field_modes
        .iter()
        .map(|(name, info)| {
            let ident = Ident::new(name, info.span);
            match info.mode {
                FieldMode::Ref => quote! { let #ident = &#receiver.#ident; },
                FieldMode::Mut => quote! { let #ident = &mut #receiver.#ident; },
                FieldMode::Copy => quote! { let #ident = #receiver.#ident; },
            }
        })
        .collect();

    // Generate final output
    let body_stmts = &body.stmts;
    Ok(quote! {
        {
            #(#bindings)*
            #(#body_stmts)*
        }
    })
}