light-sdk-macros 0.23.0

//! AST expression transformation utilities.
//!
//! This module provides expression transformation for converting field access patterns
//! used in seed derivation code generation.

use std::collections::HashSet;

use syn::Expr;

use crate::light_pdas::shared_utils::is_base_path;

// =============================================================================
// EXPRESSION TRANSFORMER
// =============================================================================

/// Transform expressions by replacing field access patterns.
///
/// Used for converting:
/// - `data.field` -> `self.field` (only if field exists in state_field_names)
/// - `ctx.field` -> `ctx_seeds.field` (if field is in ctx_field_names)
/// - `ctx.accounts.field` -> `ctx_seeds.field`
///
/// For `data.field` where field is NOT in state_field_names, the expression
/// is left unchanged (which will cause a compile error, alerting the user
/// that this field is not supported for decompression seed derivation).
pub fn transform_expr_for_ctx_seeds(
    expr: &Expr,
    ctx_field_names: &HashSet<String>,
    state_field_names: &HashSet<String>,
) -> Expr {
    transform_expr_internal(expr, ctx_field_names, state_field_names)
}

fn transform_expr_internal(
    expr: &Expr,
    ctx_field_names: &HashSet<String>,
    state_field_names: &HashSet<String>,
) -> Expr {
    match expr {
        Expr::Field(field_expr) => {
            let Some(syn::Member::Named(field_name)) = Some(&field_expr.member) else {
                return expr.clone();
            };

            // Check for ctx.accounts.field -> ctx_seeds.field
            if let Expr::Field(nested_field) = &*field_expr.base {
                if let syn::Member::Named(base_name) = &nested_field.member {
                    if base_name == "accounts" && is_base_path(&nested_field.base, "ctx") {
                        return syn::parse_quote! { ctx_seeds.#field_name };
                    }
                }
            }

            // Check for data.field -> self.field (only if field exists on state struct)
            if is_base_path(&field_expr.base, "data") {
                let field_str = field_name.to_string();
                if state_field_names.contains(&field_str) {
                    return syn::parse_quote! { self.#field_name };
                }
                // Field not on state struct - use seed_params (params-only field).
                // seed_params.field is Option<T> where T: Copy, so .unwrap() copies the value.
                return syn::parse_quote! { seed_params.#field_name.unwrap() };
            }

            // Check for ctx.field -> ctx_seeds.field
            if is_base_path(&field_expr.base, "ctx")
                && ctx_field_names.contains(&field_name.to_string())
            {
                return syn::parse_quote! { ctx_seeds.#field_name };
            }

            expr.clone()
        }
        Expr::MethodCall(method_call) => {
            let mut new_call = method_call.clone();
            new_call.receiver = Box::new(transform_expr_internal(
                &method_call.receiver,
                ctx_field_names,
                state_field_names,
            ));
            new_call.args = method_call
                .args
                .iter()
                .map(|a| transform_expr_internal(a, ctx_field_names, state_field_names))
                .collect();
            Expr::MethodCall(new_call)
        }
        Expr::Call(call_expr) => {
            let mut new_call = call_expr.clone();
            new_call.args = call_expr
                .args
                .iter()
                .map(|a| transform_expr_internal(a, ctx_field_names, state_field_names))
                .collect();
            Expr::Call(new_call)
        }
        Expr::Reference(ref_expr) => {
            let mut new_ref = ref_expr.clone();
            new_ref.expr = Box::new(transform_expr_internal(
                &ref_expr.expr,
                ctx_field_names,
                state_field_names,
            ));
            Expr::Reference(new_ref)
        }
        _ => expr.clone(),
    }
}