hamelin_translation 0.4.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.
//!
//! Expansion transforms a struct expression to match a target struct type by:
//! 1. Adding missing fields as typed NULLs
//! 2. Reordering fields to match target type's order
//! 3. Recursively expanding nested struct fields
//!
//! This module produces AST expressions using builders. Type-checking happens
//! after the pass re-typechecks the transformed pipeline/statement.
//!
//! When a complex expression needs to be accessed multiple times (to avoid
//! duplicate evaluation), hoisting is used: LET commands are returned alongside
//! the expanded expression.

use std::sync::Arc;

use hamelin_lib::tree::{
    ast::{command::Command, expression::Expression},
    builder::{
        call, cast, column_ref, drop_command, field, lambda1, let_command, null, struct_literal,
        ExpressionBuilder,
    },
    typed_ast::{
        environment::TypeEnvironment,
        expression::{TypedExpression, TypedExpressionKind, TypedStructLiteral},
    },
};
use hamelin_lib::types::{struct_type::Struct, Type};

use super::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_column_reference_chain(expr: &TypedExpression) -> bool {
    match &expr.kind {
        TypedExpressionKind::ColumnReference(_) => true,
        TypedExpressionKind::FieldLookup(lookup) => is_column_reference_chain(&lookup.value),
        _ => false,
    }
}

/// Count how many field accesses we'd need to build the expanded struct.
///
/// This counts fields that exist in source and need to be copied to target.
/// Used to determine if hoisting is necessary.
pub fn count_needed_field_accesses(source: &Struct, target: &Struct) -> usize {
    let mut count = 0;
    for (field_name, target_field_type) in target.fields.iter() {
        if let Some(source_field_type) = source.fields.get(field_name) {
            count += 1;
            // Recurse for nested structs
            if let (Type::Struct(s), Type::Struct(t)) = (source_field_type, target_field_type) {
                count += count_needed_field_accesses(s, t);
            }
        }
    }
    count
}

// ---------------------------------------------------------------------------
// 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 built with builders
/// - `before_commands` are LET 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.
pub fn expand_struct_to_type(
    expr: &Arc<TypedExpression>,
    source_type: &Struct,
    target_type: &Struct,
    name_gen: &mut UniqueNameGenerator,
    schema: &TypeEnvironment,
) -> (Arc<Expression>, Vec<Command>, Vec<Command>) {
    // Check if expansion is needed
    if source_type == target_type {
        // No expansion needed - return the original AST (cheap Rc clone)
        return (expr.ast.clone(), Vec::new(), Vec::new());
    }

    // Determine expansion strategy based on expression kind
    match &expr.kind {
        // Case 1: Struct literal - modify in place
        TypedExpressionKind::StructLiteral(lit) => {
            let (expanded, before, after) =
                expand_struct_literal(lit, target_type, name_gen, schema);
            (Arc::new(expanded), before, after)
        }

        // Case 2: Column reference chain - can duplicate freely
        _ if is_column_reference_chain(expr) => (
            Arc::new(build_expanded_struct(expr, source_type, target_type)),
            Vec::new(),
            Vec::new(),
        ),

        // Case 3: Complex expression - may need hoisting
        _ => {
            let field_access_count = count_needed_field_accesses(source_type, target_type);
            if field_access_count > 1 {
                // Hoist the expression to avoid duplicate evaluation
                let hoisted_name = name_gen.next(schema);
                let let_cmd = let_command()
                    .named_field(hoisted_name.clone(), expr.ast.clone())
                    .build();
                let drop_cmd = drop_command().field(hoisted_name.clone()).build();
                let expanded = build_expanded_struct_fields_from_column(
                    hoisted_name.as_str(),
                    source_type,
                    target_type,
                );
                (Arc::new(expanded), vec![let_cmd], vec![drop_cmd])
            } else if field_access_count == 1 {
                // Single field access, no duplication needed
                (
                    Arc::new(build_expanded_struct(expr, source_type, target_type)),
                    Vec::new(),
                    Vec::new(),
                )
            } else {
                // No field accesses needed (all fields are NULL?) - rare case
                (
                    Arc::new(cast(null(), target_type.clone().into()).build()),
                    Vec::new(),
                    Vec::new(),
                )
            }
        }
    }
}

/// Expand a struct literal to match target type.
///
/// Modifies fields in place where possible, recurses for nested expansion.
fn expand_struct_literal(
    lit: &TypedStructLiteral,
    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.fields.iter() {
        // Try to find this field in the literal
        let existing = lit.fields.iter().find(|(n, _)| {
            n.valid_ref()
                .map(|s| s.as_str() == field_name.name.as_str())
                .unwrap_or(false)
        });

        if let Some((_, field_expr)) = existing {
            // 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() {
                    // Recursively expand
                    let (expanded, before, after) = expand_struct_to_type(
                        field_expr,
                        source_nested,
                        target_nested,
                        name_gen,
                        schema,
                    );
                    all_before.extend(before);
                    all_after.extend(after);
                    builder = builder.field(field_name.clone(), 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 {
                            // Need to transform array elements: transform(field, __item -> {...})
                            let lambda_body = build_expanded_struct_fields(
                                column_ref("__item").build(),
                                source_elem,
                                target_elem,
                            );
                            let transformed = call("transform")
                                .arg(field_expr.ast.clone())
                                .arg(lambda1("__item").body(lambda_body))
                                .build();
                            builder = builder.field(field_name.clone(), Arc::new(transformed));
                            continue;
                        }
                    }
                }
            }

            // No nested expansion, keep original AST
            builder = builder.field(field_name.clone(), field_expr.ast.clone());
        } else {
            // Field missing - insert typed NULL
            builder = builder.field(field_name.clone(), cast(null(), field_type.clone()));
        }
    }

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

/// Build a struct literal that extracts fields from a typed expression.
///
/// Used when the source is a column reference or has been hoisted.
fn build_expanded_struct(
    source_expr: &Arc<TypedExpression>,
    source_type: &Struct,
    target_type: &Struct,
) -> Expression {
    let mut builder = struct_literal();

    // Build fields in target type's order
    for (field_name, field_type) in target_type.fields.iter() {
        if let Some(source_field_type) = source_type.fields.get(field_name) {
            // Field exists in source - create field access on original AST
            let field_access = field(source_expr.ast.clone(), field_name.name.as_str());

            // Check if nested struct expansion needed
            if let (Type::Struct(source_nested), Type::Struct(target_nested)) =
                (source_field_type, field_type)
            {
                if source_nested != target_nested {
                    let expanded = build_expanded_struct_fields(
                        field_access.build(),
                        source_nested,
                        target_nested,
                    );
                    builder = builder.field(field_name.clone(), expanded);
                    continue;
                }
            }

            // Check if nested array of structs expansion needed
            if let (Type::Array(source_arr), Type::Array(target_arr)) =
                (source_field_type, field_type)
            {
                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 {
                        // Need to transform array elements: transform(field, __item -> {...})
                        let lambda_body = build_expanded_struct_fields_from_column(
                            "__item",
                            source_elem,
                            target_elem,
                        );
                        let transformed = call("transform")
                            .arg(field_access)
                            .arg(lambda1("__item").body(lambda_body))
                            .build();
                        builder = builder.field(field_name.clone(), transformed);
                        continue;
                    }
                }
            }

            builder = builder.field(field_name.clone(), field_access);
        } else {
            // Field doesn't exist in source - insert typed NULL
            builder = builder.field(field_name.clone(), cast(null(), field_type.clone()));
        }
    }

    builder.build()
}

/// 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 `column_ref(field_name)` or `cast(null(), type)`
/// For nested structs: recursively builds struct literal with widened fields
/// For arrays of structs: wraps in `transform(array, x -> {...})` to widen elements
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 nested struct expansion needed
            if let (Type::Struct(source_struct), Type::Struct(target_struct)) = (source_t, target_t)
            {
                if source_struct != target_struct {
                    // Need to expand nested struct
                    return build_expanded_struct_fields(
                        column_ref(field_name).build(),
                        source_struct,
                        target_struct,
                    );
                }
            }
            // 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 {
                        // Need to transform array elements: transform(field, __item -> {...})
                        let lambda_body = build_expanded_struct_fields_from_column(
                            "__item",
                            source_elem,
                            target_elem,
                        );
                        return call("transform")
                            .arg(column_ref(field_name))
                            .arg(lambda1("__item").body(lambda_body))
                            .build();
                    }
                }
            }
            // Simple field reference
            column_ref(field_name).build()
        }
        // Field doesn't exist in source - use typed NULL
        (None, target_t) => cast(null(), target_t.clone()).build(),
    }
}

/// Build expanded struct fields from an expression (for nested expansion).
///
/// This is used when we don't have a TypedExpression to work with (e.g., when
/// building nested field accesses).
fn build_expanded_struct_fields(
    source_expr: Expression,
    source_type: &Struct,
    target_type: &Struct,
) -> Expression {
    let mut builder = struct_literal();

    for (field_name, field_type) in target_type.fields.iter() {
        if let Some(source_field_type) = source_type.fields.get(field_name) {
            let field_access = field(source_expr.clone(), field_name.name.as_str());

            // Check if nested struct expansion needed
            if let (Type::Struct(source_nested), Type::Struct(target_nested)) =
                (source_field_type, field_type)
            {
                if source_nested != target_nested {
                    let expanded = build_expanded_struct_fields(
                        field_access.build(),
                        source_nested,
                        target_nested,
                    );
                    builder = builder.field(field_name.clone(), expanded);
                    continue;
                }
            }

            // Check if nested array of structs expansion needed
            if let (Type::Array(source_arr), Type::Array(target_arr)) =
                (source_field_type, field_type)
            {
                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 {
                        // Need to transform array elements: transform(field, __item -> {...})
                        let lambda_body = build_expanded_struct_fields_from_column(
                            "__item",
                            source_elem,
                            target_elem,
                        );
                        let transformed = call("transform")
                            .arg(field_access)
                            .arg(lambda1("__item").body(lambda_body))
                            .build();
                        builder = builder.field(field_name.clone(), transformed);
                        continue;
                    }
                }
            }

            builder = builder.field(field_name.clone(), field_access);
        } else {
            builder = builder.field(field_name.clone(), cast(null(), field_type.clone()));
        }
    }

    builder.build()
}

/// Build expanded struct fields from a column reference name.
///
/// Similar to `build_expanded_struct_fields` but takes a column name instead of
/// an expression. Used for building lambda bodies in array transforms.
fn build_expanded_struct_fields_from_column(
    column_name: &str,
    source_type: &Struct,
    target_type: &Struct,
) -> Expression {
    build_expanded_struct_fields(column_ref(column_name).build(), source_type, target_type)
}

#[cfg(test)]
mod tests {
    use super::*;
    use hamelin_lib::tree::ast::expression::Expression;
    use hamelin_lib::tree::ast::{IntoTyped, TypeCheckExecutor};
    use hamelin_lib::tree::builder::{
        call, cast, column_ref, field, ident, lambda1, null, struct_literal,
    };
    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 count_needed_field_accesses
    // -------------------------------------------------------------------------

    #[rstest]
    #[case::flat_struct(
        Struct::default().with_str("a", INT).with_str("b", INT),
        Struct::default().with_str("a", INT).with_str("b", INT).with_str("c", INT),
        2  // a and b exist in source, c doesn't
    )]
    #[case::nested_struct(
        Struct::default().with_str("nested", Struct::default().with_str("x", INT).into()),
        Struct::default().with_str("nested", Struct::default().with_str("x", INT).with_str("y", INT).into()),
        2  // 1 for nested field + 1 for x inside nested
    )]
    fn test_count_needed_field_accesses(
        #[case] source: Struct,
        #[case] target: Struct,
        #[case] expected: usize,
    ) {
        assert_eq!(count_needed_field_accesses(&source, &target), expected);
    }

    // -------------------------------------------------------------------------
    // 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,
        column_ref("existing_field").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(),
        call("transform")
            .arg(column_ref("items"))
            .arg(lambda1("__item").body(
                struct_literal()
                    .field("a", field(column_ref("__item"), "a"))
                    .field("b", cast(null(), INT)),
            ))
            .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 build_expanded_struct_fields
    // -------------------------------------------------------------------------

    #[rstest]
    #[case::reorders_fields(
        Struct::default().with_str("b", INT).with_str("a", INT),
        Struct::default().with_str("a", INT).with_str("b", INT).with_str("c", INT),
        struct_literal()
            .field("a", field(column_ref("source"), "a"))
            .field("b", field(column_ref("source"), "b"))
            .field("c", cast(null(), INT))
            .build()
    )]
    #[case::nested_struct(
        Struct::default().with_str("nested", Struct::default().with_str("x", INT).into()),
        Struct::default().with_str("nested", Struct::default().with_str("x", INT).with_str("y", INT).into()),
        struct_literal()
            .field(
                "nested",
                struct_literal()
                    .field("x", field(field(column_ref("source"), "nested"), "x"))
                    .field("y", cast(null(), INT)),
            )
            .build()
    )]
    #[case::nested_array(
        Struct::default().with_str("items", Array::new(Struct::default().with_str("a", INT).into()).into()),
        Struct::default().with_str("items", Array::new(Struct::default().with_str("a", INT).with_str("b", INT).into()).into()),
        struct_literal()
            .field(
                "items",
                call("transform")
                    .arg(field(column_ref("source"), "items"))
                    .arg(lambda1("__item").body(
                        struct_literal()
                            .field("a", field(column_ref("__item"), "a"))
                            .field("b", cast(null(), INT)),
                    )),
            )
            .build()
    )]
    fn test_build_expanded_struct_fields(
        #[case] source: Struct,
        #[case] target: Struct,
        #[case] expected: Expression,
    ) {
        let result = build_expanded_struct_fields(column_ref("source").build(), &source, &target);
        assert_eq!(result, expected);
    }

    // -------------------------------------------------------------------------
    // Tests for is_column_reference_chain
    // -------------------------------------------------------------------------

    fn type_check_expr(expr: Expression, bindings: Arc<TypeEnvironment>) -> TypedExpression {
        expr.typed_with().with_bindings(bindings).typed()
    }

    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_base(ident("col").into(), INT)
                .unwrap()
                .with_base(ident("s").into(), outer_struct)
                .unwrap(),
        )
    }

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

    // -------------------------------------------------------------------------
    // Tests for expand_struct_to_type (array-of-structs via column-ref/hoisted paths)
    // -------------------------------------------------------------------------

    #[rstest]
    #[case::column_ref_array_of_structs(
        // Source: {items: Array<{a: INT}>}, column reference path
        column_ref("data").build(),
        Struct::default().with_str("items", Array::new(Struct::default().with_str("a", INT).into()).into()),
        Struct::default().with_str("items", Array::new(Struct::default().with_str("a", INT).with_str("b", INT).into()).into()),
        0,  // no hoisting for column refs
        struct_literal()
            .field(
                "items",
                call("transform")
                    .arg(field(column_ref("data"), "items"))
                    .arg(lambda1("__item").body(
                        struct_literal()
                            .field("a", field(column_ref("__item"), "a"))
                            .field("b", cast(null(), INT)),
                    )),
            )
            .build()
    )]
    #[case::hoisted_array_of_structs(
        // Source: {x: INT, items: Array<{a: INT}>}, complex expr triggers hoisting
        call("coalesce").arg(column_ref("data")).arg(column_ref("data")).build(),
        Struct::default().with_str("x", INT).with_str("items", Array::new(Struct::default().with_str("a", INT).into()).into()),
        Struct::default().with_str("x", INT).with_str("items", Array::new(Struct::default().with_str("a", INT).with_str("b", INT).into()).into()),
        1,  // hoisting needed for complex expr with 2+ field accesses
        struct_literal()
            .field("x", field(column_ref("__test_0"), "x"))
            .field(
                "items",
                call("transform")
                    .arg(field(column_ref("__test_0"), "items"))
                    .arg(lambda1("__item").body(
                        struct_literal()
                            .field("a", field(column_ref("__item"), "a"))
                            .field("b", cast(null(), INT)),
                    )),
            )
            .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_base(ident("data").into(), source_type.clone().into())
                .unwrap(),
        );
        let typed_expr = Arc::new(expr.typed_with().with_bindings(bindings.clone()).typed());

        let mut name_gen = UniqueNameGenerator::new("__test");
        let (result, before, after) = expand_struct_to_type(
            &typed_expr,
            &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);
    }
}