depyler-core 3.18.1

//! Generator support and code generation
//!
//! This module handles Python generator functions, converting them to
//! Rust Iterator implementations with state structs.

use crate::generator_state::GeneratorStateInfo;
use crate::hir::{HirFunction, Type};
use crate::rust_gen::CodeGenContext;
use crate::rust_gen::type_gen::rust_type_to_syn;
use anyhow::Result;
use quote::quote;

/// Generate struct fields for generator state variables
///
/// Creates field declarations for variables that need to persist across yields.
///
/// # Complexity
/// 3 (iter + map + collect)
fn generate_state_fields(
    state_info: &GeneratorStateInfo,
    ctx: &mut CodeGenContext,
) -> Result<Vec<proc_macro2::TokenStream>> {
    state_info
        .state_variables
        .iter()
        .map(|var| {
            let field_name = syn::Ident::new(&var.name, proc_macro2::Span::call_site());
            let rust_type = ctx.type_mapper.map_type(&var.ty);
            let field_type = rust_type_to_syn(&rust_type)?;
            Ok(quote! { #field_name: #field_type })
        })
        .collect()
}

/// Generate struct fields for captured parameters
///
/// Creates field declarations for function parameters that are captured
/// in the generator state.
///
/// # Complexity
/// 4 (iter + filter + map + collect)
fn generate_param_fields(
    func: &HirFunction,
    state_info: &GeneratorStateInfo,
    ctx: &mut CodeGenContext,
) -> Result<Vec<proc_macro2::TokenStream>> {
    func.params
        .iter()
        .filter(|p| state_info.captured_params.contains(&p.name))
        .map(|param| {
            let field_name = syn::Ident::new(&param.name, proc_macro2::Span::call_site());
            let rust_type = ctx.type_mapper.map_type(&param.ty);
            let field_type = rust_type_to_syn(&rust_type)?;
            Ok(quote! { #field_name: #field_type })
        })
        .collect()
}

/// Extract the Item type for the Iterator from generator return type
///
/// Generator expressions are fully implemented in v3.13.0 (20/20 tests passing).
/// This function uses the return type directly as the Iterator::Item type.
///
/// # Complexity
/// 1 (simple delegation)
#[inline]
fn extract_generator_item_type(
    rust_ret_type: &crate::type_mapper::RustType,
) -> Result<syn::Type> {
    rust_type_to_syn(rust_ret_type)
}

/// Generate field initializers for state variables (with default values)
///
/// Creates initialization expressions like `field_name: 0` or `field_name: false`.
///
/// # Complexity
/// 3 (iter + map + collect)
fn generate_state_initializers(
    state_info: &GeneratorStateInfo,
) -> Vec<proc_macro2::TokenStream> {
    state_info
        .state_variables
        .iter()
        .map(|var| {
            let field_name = syn::Ident::new(&var.name, proc_macro2::Span::call_site());
            // Initialize with type-appropriate default (0 for int, false for bool, etc.)
            let default_value = get_default_value_for_type(&var.ty);
            quote! { #field_name: #default_value }
        })
        .collect()
}

/// Generate field initializers for captured parameters
///
/// Creates initialization expressions like `field_name: field_name` to capture
/// the parameter value.
///
/// # Complexity
/// 4 (iter + filter + map + collect)
fn generate_param_initializers(
    func: &HirFunction,
    state_info: &GeneratorStateInfo,
) -> Vec<proc_macro2::TokenStream> {
    func.params
        .iter()
        .filter(|p| state_info.captured_params.contains(&p.name))
        .map(|param| {
            let field_name = syn::Ident::new(&param.name, proc_macro2::Span::call_site());
            // Initialize with parameter value (n: n)
            quote! { #field_name: #field_name }
        })
        .collect()
}

/// Get default value expression for Int type
///
/// # Complexity: 1
#[inline]
fn default_int() -> proc_macro2::TokenStream {
    quote! { 0 }
}

/// Get default value expression for Float type
///
/// # Complexity: 1
#[inline]
fn default_float() -> proc_macro2::TokenStream {
    quote! { 0.0 }
}

/// Get default value expression for Bool type
///
/// # Complexity: 1
#[inline]
fn default_bool() -> proc_macro2::TokenStream {
    quote! { false }
}

/// Get default value expression for String type
///
/// # Complexity: 1
#[inline]
fn default_string() -> proc_macro2::TokenStream {
    quote! { String::new() }
}

/// Get default value expression for other types
///
/// # Complexity: 1
#[inline]
fn default_generic() -> proc_macro2::TokenStream {
    quote! { Default::default() }
}

/// Get default value expression for a type
///
/// Returns appropriate default value based on HIR type.
///
/// # Complexity
/// 6 (match with 5 arms)
#[inline]
fn get_default_value_for_type(ty: &Type) -> proc_macro2::TokenStream {
    match ty {
        Type::Int => default_int(),
        Type::Float => default_float(),
        Type::Bool => default_bool(),
        Type::String => default_string(),
        _ => default_generic(),
    }
}

/// Generate state struct name by capitalizing function name
///
/// # Complexity: 2
#[inline]
fn generate_state_struct_name(name: &syn::Ident) -> syn::Ident {
    let state_struct_name = format!(
        "{}State",
        name.to_string()
            .chars()
            .next()
            .map(|c| c.to_uppercase().to_string())
            .unwrap_or_default()
            + &name.to_string()[1..]
    );
    syn::Ident::new(&state_struct_name, name.span())
}

/// Populate generator state variables in context
///
/// # Complexity: 3 (clear + 2 loops)
#[inline]
fn populate_generator_state_vars(
    ctx: &mut CodeGenContext,
    state_info: &GeneratorStateInfo,
) {
    ctx.generator_state_vars.clear();
    for var in &state_info.state_variables {
        ctx.generator_state_vars.insert(var.name.clone());
    }
    for param in &state_info.captured_params {
        ctx.generator_state_vars.insert(param.clone());
    }
}

/// Generate generator body statements with proper context flags
///
/// # Complexity: 4 (set flag + collect + clear flag + clear vars)
#[inline]
fn generate_generator_body(
    func: &HirFunction,
    ctx: &mut CodeGenContext,
) -> Result<Vec<proc_macro2::TokenStream>> {
    use crate::rust_gen::RustCodeGen;

    ctx.in_generator = true;
    let generator_body_stmts: Vec<_> = func
        .body
        .iter()
        .map(|stmt| stmt.to_rust_tokens(ctx))
        .collect::<Result<Vec<_>>>()?;
    ctx.in_generator = false;
    ctx.generator_state_vars.clear();

    Ok(generator_body_stmts)
}

/// Generate complete generator function with state struct and Iterator impl
///
/// This is the main entry point for generator code generation. It:
/// 1. Analyzes generator state requirements
/// 2. Creates a state struct with captured variables
/// 3. Generates a constructor function
/// 4. Implements Iterator with state machine logic
///
/// # Arguments
/// * `func` - The HIR function to generate
/// * `name` - Function name identifier
/// * `generic_params` - Generic parameters token stream
/// * `where_clause` - Where clause token stream
/// * `params` - Parameter declarations
/// * `attrs` - Function attributes
/// * `rust_ret_type` - Return type
/// * `ctx` - Code generation context
///
/// # Returns
/// Complete generator implementation including state struct and Iterator impl
///
/// # Complexity
/// 5 (delegated to helper functions)
#[inline]
#[allow(clippy::too_many_arguments)] // Generator needs all metadata for complex transformation
pub fn codegen_generator_function(
    func: &HirFunction,
    name: &syn::Ident,
    generic_params: &proc_macro2::TokenStream,
    where_clause: &proc_macro2::TokenStream,
    params: &[proc_macro2::TokenStream],
    attrs: &[proc_macro2::TokenStream],
    rust_ret_type: &crate::type_mapper::RustType,
    ctx: &mut CodeGenContext,
) -> Result<proc_macro2::TokenStream> {
    // Analyze generator state requirements
    let state_info = GeneratorStateInfo::analyze(func);

    // Generate state struct name
    let state_ident = generate_state_struct_name(name);

    // Build state struct fields from analysis
    let state_fields = generate_state_fields(&state_info, ctx)?;
    let param_fields = generate_param_fields(func, &state_info, ctx)?;
    let all_fields = [state_fields, param_fields].concat();

    // Build field initializers
    let state_inits = generate_state_initializers(&state_info);
    let param_inits = generate_param_initializers(func, &state_info);
    let all_inits = [state_inits, param_inits].concat();

    // Generate state machine field
    let state_machine_field = quote! {
        state: usize
    };

    // Extract yield value type from return type
    let item_type = extract_generator_item_type(rust_ret_type)?;

    // Populate generator state variables for scoping
    populate_generator_state_vars(ctx, &state_info);

    // Generate body statements with proper context
    let generator_body_stmts = generate_generator_body(func, ctx)?;

    // Generate the complete generator implementation
    Ok(quote! {
        #(#attrs)*
        #[doc = " Generator state struct"]
        #[derive(Debug)]
        struct #state_ident {
            #state_machine_field,
            #(#all_fields),*
        }

        #[doc = " Generator function - returns Iterator"]
        pub fn #name #generic_params(#(#params),*) -> impl Iterator<Item = #item_type> #where_clause {
            #state_ident {
                state: 0,
                #(#all_inits),*
            }
        }

        impl Iterator for #state_ident {
            type Item = #item_type;

            fn next(&mut self) -> Option<Self::Item> {
                // State machine implementation: Simplified single-state execution
                // for basic generator support (v3.12.0). Multi-state transformation
                // with resumable yield points is future work.
                match self.state {
                    0 => {
                        self.state = 1;
                        // Execute generator body with early-exit semantics
                        #(#generator_body_stmts)*
                        None
                    }
                    _ => None
                }
            }
        }
    })
}