hamelin_translation 0.9.2

//! Struct expansion helpers.
//!
//! Provides functions for expanding struct expressions to match a target struct type.
//! Used by `expand_array_literals` for array element widening.
//!
//! This module produces AST Cast expressions that widen structs to target types.
//! The actual expansion (adding NULL fields, reordering) is handled by the translation
//! layer (SQL or DataFusion) based on `CastKind::StructExpansion`.
//!
//! When a complex expression needs to be accessed multiple times (to avoid
//! duplicate evaluation), hoisting is used: SET commands are returned alongside
//! the expanded expression.

use std::sync::Arc;

use hamelin_lib::tree::{
    ast::{command::Command, expression::Expression, identifier::ParsedSimpleIdentifier},
    builder::{cast, drop_command, field_ref, set_command, struct_literal, ExpressionBuilder},
    typed_ast::{
        environment::TypeEnvironment,
        expression::{TypedExpression, TypedExpressionKind, TypedStructLiteral},
    },
};
#[cfg(test)]
use hamelin_lib::type_check_expression;
use hamelin_lib::types::{array::Array, struct_type::Struct, Type};

use crate::unique::UniqueNameGenerator;

// ---------------------------------------------------------------------------
// Helper Functions
// ---------------------------------------------------------------------------

/// Check if an expression is a column reference chain (col, col.field, col.field.nested).
///
/// These can be duplicated without performance penalty since they're just
/// lookups, not computations.
pub fn is_field_reference_chain(expr: &TypedExpression) -> bool {
    match &expr.kind {
        TypedExpressionKind::FieldReference(_) => true,
        TypedExpressionKind::FieldLookup(lookup) => is_field_reference_chain(&lookup.value),
        _ => false,
    }
}

/// Check if an expression is simple enough to duplicate without concern.
fn is_simple_expression(expr: &TypedExpression) -> bool {
    match &expr.kind {
        TypedExpressionKind::FieldReference(_) => true,
        TypedExpressionKind::FieldLookup(lookup) => is_field_reference_chain(&lookup.value),
        // Leaf covers all simple literals (int, float, string, bool, null, etc.)
        TypedExpressionKind::Leaf => true,
        TypedExpressionKind::Cast(c) => is_simple_expression(&c.value),
        // Struct literals are simple if all fields are simple
        TypedExpressionKind::StructLiteral(lit) => lit
            .fields
            .iter()
            .all(|(_, field_expr)| is_simple_expression(field_expr)),
        // Array literals are simple if all elements are simple
        TypedExpressionKind::ArrayLiteral(arr) => {
            arr.elements.iter().all(|elem| is_simple_expression(elem))
        }
        _ => false,
    }
}

// ---------------------------------------------------------------------------
// Core Expansion Functions
// ---------------------------------------------------------------------------

/// Expand a typed struct expression to match a target struct type.
///
/// Returns `(expanded_expression, before_commands, after_commands)` where:
/// - `expanded_expression` is the AST expression (either a Cast or modified struct literal)
/// - `before_commands` are SET commands that must be inserted before the command
///   that uses this expression (to avoid duplicate evaluation of complex expressions)
/// - `after_commands` are DROP commands that must be inserted after the command
///   to clean up the hoisted variables
///
/// The caller is responsible for:
/// - Ensuring `expr` has a struct type matching `source_type`
/// - Re-typechecking after all transformations are complete
///
/// The `schema` parameter is used to ensure generated names don't collide with
/// existing columns.
/// Expand a typed struct expression to match a target struct type.
///
/// When `transformed_ast` is provided, it is used in place of `expr.ast` to
/// preserve any nested rewrites already applied by a bottom-up traversal
/// (e.g., inner array literal expansions from a cata pass).
pub fn expand_struct_to_type_with_ast(
    expr: &Arc<TypedExpression>,
    transformed_ast: Option<&Arc<Expression>>,
    source_type: &Struct,
    target_type: &Struct,
    name_gen: &mut UniqueNameGenerator,
    schema: &TypeEnvironment,
) -> (Arc<Expression>, Vec<Command>, Vec<Command>) {
    let ast = transformed_ast.unwrap_or(&expr.ast);

    // Check if expansion is needed
    if source_type == target_type {
        // No expansion needed - return the (possibly transformed) AST
        return (ast.clone(), Vec::new(), Vec::new());
    }

    // Determine expansion strategy based on expression kind
    match &expr.kind {
        // Case 1: Struct literal - expand inline by rebuilding with added fields
        TypedExpressionKind::StructLiteral(lit) => {
            // Extract transformed field values from the transformed AST if available
            let transformed_fields = transformed_ast.and_then(|a| match &a.kind {
                hamelin_lib::tree::ast::expression::ExpressionKind::StructLiteral(s) => {
                    Some(&s.fields)
                }
                _ => None,
            });
            let (expanded, before, after) =
                expand_struct_literal(lit, transformed_fields, target_type, name_gen, schema);
            (Arc::new(expanded), before, after)
        }

        // Case 2: Simple expression (column reference, literals, etc.) - just wrap in Cast
        _ if is_field_reference_chain(expr) || is_simple_expression(expr) => {
            let cast_expr = cast(ast.clone(), target_type.clone().into()).build();
            (Arc::new(cast_expr), Vec::new(), Vec::new())
        }

        // Case 3: Complex expression - hoist first, then cast the hoisted variable
        _ => {
            // Hoist the expression to avoid duplicate evaluation
            let hoisted_name = name_gen.next(schema);
            let set_cmd = set_command()
                .named_field(hoisted_name.clone(), ast.clone())
                .build();
            let drop_cmd = drop_command().field(hoisted_name.clone()).build();

            // Cast the hoisted column reference
            let cast_expr =
                cast(field_ref(hoisted_name.as_str()), target_type.clone().into()).build();

            (Arc::new(cast_expr), vec![set_cmd], vec![drop_cmd])
        }
    }
}

/// Expand a struct literal to match target type.
///
/// Modifies fields in place where possible, recurses for nested expansion.
fn expand_struct_literal(
    lit: &TypedStructLiteral,
    transformed_fields: Option<&Vec<(ParsedSimpleIdentifier, Arc<Expression>)>>,
    target_type: &Struct,
    name_gen: &mut UniqueNameGenerator,
    schema: &TypeEnvironment,
) -> (Expression, Vec<Command>, Vec<Command>) {
    let mut builder = struct_literal();
    let mut all_before = Vec::new();
    let mut all_after = Vec::new();

    // Build fields in target type's order
    for (field_name, field_type) in target_type.iter() {
        // Try to find this field in the literal (returns index for transformed lookup)
        let existing = lit.fields.iter().enumerate().find(|(_, (n, _))| {
            n.valid_ref()
                .map(|s| s.as_str() == field_name.name())
                .unwrap_or(false)
        });

        if let Some((idx, (_, field_expr))) = existing {
            // Get the transformed AST for this field if available
            let transformed_field_ast = transformed_fields.map(|tf| &tf[idx].1);

            // Field exists in literal
            // Check if nested struct expansion needed
            if let Type::Struct(target_nested) = field_type {
                if let Type::Struct(source_nested) = field_expr.resolved_type.as_ref() {
                    if source_nested != target_nested {
                        // Recursively expand nested struct, passing transformed AST
                        let (expanded, before, after) = expand_struct_to_type_with_ast(
                            field_expr,
                            transformed_field_ast,
                            source_nested,
                            target_nested,
                            name_gen,
                            schema,
                        );
                        all_before.extend(before);
                        all_after.extend(after);
                        builder = builder.field(field_name.name(), expanded);
                        continue;
                    }
                }
            }

            // Check if nested array of structs expansion needed
            if let Type::Array(target_arr) = field_type {
                if let Type::Array(source_arr) = field_expr.resolved_type.as_ref() {
                    if let (Type::Struct(source_elem), Type::Struct(target_elem)) = (
                        source_arr.element_type.as_ref(),
                        target_arr.element_type.as_ref(),
                    ) {
                        if source_elem != target_elem {
                            // Cast the array to the target array type, using transformed AST
                            let field_ast =
                                transformed_field_ast.unwrap_or(&field_expr.ast).clone();
                            let target_array_type: Type =
                                Array::new(target_elem.clone().into()).into();
                            let cast_expr = cast(field_ast, target_array_type).build();
                            builder = builder.field(field_name.name(), cast_expr);
                            continue;
                        }
                    }
                }
            }

            // No nested expansion — use transformed AST to preserve nested rewrites
            let field_ast = transformed_field_ast.unwrap_or(&field_expr.ast).clone();
            builder = builder.field(field_name.name(), field_ast);
        } else {
            // Field missing - insert typed NULL via cast
            builder = builder.field(
                field_name.name(),
                cast(hamelin_lib::tree::builder::null(), field_type.clone()),
            );
        }
    }

    (builder.build(), all_before, all_after)
}

/// Build an expression that widens a source type to a target type.
///
/// Used by `expand_union_schemas` to build SELECT projections that widen
/// individual fields from source schema to target schema.
///
/// For scalar fields: returns `field_ref(field_name)` or `cast(null(), type)`
/// For nested structs: returns `cast(field_ref(field_name), target_struct)`
/// For arrays of structs: returns `cast(field_ref(field_name), array<target_struct>)`
pub fn build_widening_expression(
    field_name: &str,
    source_type: Option<&Type>,
    target_type: &Type,
) -> Expression {
    match (source_type, target_type) {
        // Field exists in source
        (Some(source_t), target_t) => {
            // Check if struct expansion needed
            if let (Type::Struct(source_struct), Type::Struct(target_struct)) = (source_t, target_t)
            {
                if source_struct != target_struct {
                    // Cast to target struct type - CastKind::StructExpansion will handle it
                    return cast(field_ref(field_name), target_t.clone()).build();
                }
            }

            // Check if array of structs expansion needed
            if let (Type::Array(source_arr), Type::Array(target_arr)) = (source_t, target_t) {
                if let (Type::Struct(source_elem), Type::Struct(target_elem)) = (
                    source_arr.element_type.as_ref(),
                    target_arr.element_type.as_ref(),
                ) {
                    if source_elem != target_elem {
                        // Cast to target array type - CastKind::ArrayElementCast will handle it
                        return cast(field_ref(field_name), target_t.clone()).build();
                    }
                }
            }

            // Simple field reference (no widening needed or just implicit cast)
            field_ref(field_name).build()
        }
        // Field doesn't exist in source - use typed NULL
        (None, target_t) => cast(hamelin_lib::tree::builder::null(), target_t.clone()).build(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use hamelin_lib::tree::ast::expression::Expression;
    use hamelin_lib::tree::builder::{call, cast, field_ref, null};
    use hamelin_lib::tree::options::ExpressionTypeCheckOptions;
    use hamelin_lib::tree::typed_ast::environment::TypeEnvironment;
    use hamelin_lib::tree::typed_ast::expression::TypedExpression;
    use hamelin_lib::types::array::Array;
    use hamelin_lib::types::INT;
    use pretty_assertions::assert_eq;
    use rstest::rstest;

    // -------------------------------------------------------------------------
    // Tests for build_widening_expression
    // -------------------------------------------------------------------------

    #[rstest]
    #[case::missing_field(
        "missing_field",
        None,
        INT,
        cast(null(), INT).build()
    )]
    #[case::existing_field(
        "existing_field",
        Some(INT),
        INT,
        field_ref("existing_field").build()
    )]
    #[case::nested_struct_widening(
        "nested",
        Some(Struct::default().with_str("a", INT).into()),
        Struct::default().with_str("a", INT).with_str("b", INT).into(),
        cast(
            field_ref("nested"),
            Struct::default().with_str("a", INT).with_str("b", INT).into(),
        ).build()
    )]
    #[case::array_of_structs(
        "items",
        Some(Array::new(Struct::default().with_str("a", INT).into()).into()),
        Array::new(Struct::default().with_str("a", INT).with_str("b", INT).into()).into(),
        cast(
            field_ref("items"),
            Array::new(Struct::default().with_str("a", INT).with_str("b", INT).into()).into(),
        ).build()
    )]
    fn test_build_widening_expression(
        #[case] field_name: &str,
        #[case] source_type: Option<Type>,
        #[case] target_type: Type,
        #[case] expected: Expression,
    ) {
        let result = build_widening_expression(field_name, source_type.as_ref(), &target_type);
        assert_eq!(result, expected);
    }

    // -------------------------------------------------------------------------
    // Tests for is_field_reference_chain
    // -------------------------------------------------------------------------

    fn type_check_expr(expr: Expression, bindings: Arc<TypeEnvironment>) -> TypedExpression {
        type_check_expression(
            expr,
            ExpressionTypeCheckOptions::builder()
                .bindings(bindings)
                .build(),
        )
        .output
    }

    fn test_bindings() -> Arc<TypeEnvironment> {
        let nested_struct: Type = Struct::default().with_str("inner", INT).into();
        let outer_struct: Type = Struct::default().with_str("field", nested_struct).into();
        Arc::new(
            TypeEnvironment::default()
                .with(hamelin_lib::tree::builder::ident("col").into(), INT)
                .with(hamelin_lib::tree::builder::ident("s").into(), outer_struct),
        )
    }

    #[rstest]
    #[case::simple_column(field_ref("col").build(), true)]
    #[case::one_field_access(hamelin_lib::tree::builder::field(field_ref("s"), "field").build(), true)]
    #[case::nested_field_access(hamelin_lib::tree::builder::field(hamelin_lib::tree::builder::field(field_ref("s"), "field"), "inner").build(), true)]
    #[case::function_call(call("coalesce").arg(field_ref("col")).arg(0).build(), false)]
    #[case::binary_operation(hamelin_lib::tree::builder::add(field_ref("col"), 1).build(), false)]
    fn test_is_field_reference_chain(#[case] expr: Expression, #[case] expected: bool) {
        let bindings = test_bindings();
        let typed_expr = type_check_expr(expr, bindings);
        assert_eq!(is_field_reference_chain(&typed_expr), expected);
    }

    // -------------------------------------------------------------------------
    // Tests for expand_struct_to_type
    // -------------------------------------------------------------------------

    fn ident(s: &str) -> hamelin_lib::tree::ast::identifier::SimpleIdentifier {
        hamelin_lib::tree::ast::identifier::SimpleIdentifier::new(s)
    }

    #[rstest]
    #[case::field_ref_emits_cast(
        // Source: {a: INT, b: INT}, column reference path
        field_ref("data").build(),
        Struct::default().with_str("a", INT).with_str("b", INT),
        Struct::default().with_str("a", INT).with_str("b", INT).with_str("c", INT),
        0,  // no hoisting for column refs
        cast(
            field_ref("data"),
            Struct::default().with_str("a", INT).with_str("b", INT).with_str("c", INT).into(),
        ).build()
    )]
    #[case::complex_expr_hoists_then_casts(
        // Complex expr triggers hoisting
        call("coalesce").arg(field_ref("data")).arg(field_ref("data")).build(),
        Struct::default().with_str("a", INT).with_str("b", INT),
        Struct::default().with_str("a", INT).with_str("b", INT).with_str("c", INT),
        1,  // hoisting needed for complex expr
        cast(
            field_ref("__test_0"),
            Struct::default().with_str("a", INT).with_str("b", INT).with_str("c", INT).into(),
        ).build()
    )]
    fn test_expand_struct_to_type(
        #[case] expr: Expression,
        #[case] source_type: Struct,
        #[case] target_type: Struct,
        #[case] expected_hoisted_count: usize,
        #[case] expected: Expression,
    ) {
        let bindings = Arc::new(
            TypeEnvironment::default().with(ident("data").into(), source_type.clone().into()),
        );
        let typed_expr = Arc::new(
            type_check_expression(
                expr,
                ExpressionTypeCheckOptions::builder()
                    .bindings(bindings.clone())
                    .build(),
            )
            .output,
        );

        let mut name_gen = UniqueNameGenerator::new("__test");
        let (result, before, after) = expand_struct_to_type_with_ast(
            &typed_expr,
            None,
            &source_type,
            &target_type,
            &mut name_gen,
            &bindings,
        );

        assert_eq!(before.len(), expected_hoisted_count);
        assert_eq!(after.len(), expected_hoisted_count);
        assert_eq!(*result, expected);
    }
}