hamelin_translation 0.7.3

//! Pipeline pass: Array literal expansion.
//!
//! Expands array literal elements to match the array's element type.
//! Each element that is a struct is expanded to match the array's struct element type.
//!
//! Example:
//! ```text
//! LET arr = [{a: 1}, {a: 2, b: 3}]
//! -- Element type is {a: Int, b: Int?}
//! -- First element is missing field b
//! ```
//! becomes:
//! ```text
//! LET arr = [{a: 1, b: CAST(NULL AS Int)}, {a: 2, b: 3}]
//! ```
//!
//! For non-literal struct expressions (column refs, complex exprs), the pass emits
//! explicit Cast expressions. The translation layer handles these via CastKind::StructExpansion.
//!
//! When expanding requires hoisting (to avoid duplicate evaluation of complex
//! expressions), this pass inserts LET commands before the command that uses them.

use std::sync::Arc;

use hamelin_lib::{
    err::TranslationError,
    tree::{
        ast::{command::Command, expression::Expression, pipeline::Pipeline},
        builder::{array, pipeline, ExpressionBuilder},
        typed_ast::{
            command::TypedCommand,
            context::StatementTranslationContext,
            environment::TypeEnvironment,
            expression::{
                MapExpressionAlgebra, TypedArrayLiteral, TypedExpression, TypedExpressionKind,
            },
            pipeline::TypedPipeline,
        },
    },
    types::{struct_type::Struct, Type},
};

use super::super::expand_struct::expand_struct_to_type_with_ast;
use crate::unique::UniqueNameGenerator;

// ---------------------------------------------------------------------------
// Public Entry Point
// ---------------------------------------------------------------------------

/// Expand array literals in a pipeline.
///
/// This is a pipeline-level pass that transforms array literals with struct
/// elements that don't match the array's element type.
///
/// Contract: `(Arc<TypedPipeline>, &mut ctx) -> Result<Arc<TypedPipeline>, ...>`
pub fn expand_array_literals(
    pipeline: Arc<TypedPipeline>,
    ctx: &mut StatementTranslationContext,
) -> Result<Arc<TypedPipeline>, Arc<TranslationError>> {
    // Check if any command needs array expansion
    if !pipeline_needs_expansion(&pipeline)? {
        return Ok(pipeline);
    }

    // Transform the pipeline
    let new_ast = transform_pipeline(&pipeline)?;

    // Re-typecheck
    Ok(Arc::new(TypedPipeline::from_ast_with_context(
        Arc::new(new_ast),
        ctx,
    )))
}

// ---------------------------------------------------------------------------
// Check if Expansion is Needed
// ---------------------------------------------------------------------------

/// Check if a pipeline has any array literals that need expansion.
fn pipeline_needs_expansion(pipeline: &TypedPipeline) -> Result<bool, Arc<TranslationError>> {
    let valid = pipeline.valid_ref()?;
    Ok(valid.commands.iter().any(command_needs_expansion))
}

/// Check if a command has any array literals that need expansion.
fn command_needs_expansion(cmd: &Arc<TypedCommand>) -> bool {
    cmd.find_expression(&mut |expr| expression_needs_expansion(expr))
        .is_some()
}

/// Check if an expression (or any descendant) contains arrays needing expansion.
fn expression_needs_expansion(expr: &TypedExpression) -> bool {
    expr.find(&mut |e| {
        let TypedExpressionKind::ArrayLiteral(arr) = &e.kind else {
            return false;
        };
        let Type::Array(arr_type) = e.resolved_type.as_ref() else {
            return false;
        };
        let Type::Struct(element_type) = arr_type.element_type.as_ref() else {
            return false;
        };
        arr.elements
            .iter()
            .any(|elem| matches!(elem.resolved_type.as_ref(), Type::Struct(s) if s != element_type))
    })
    .is_some()
}

// ---------------------------------------------------------------------------
// Transform Pipeline
// ---------------------------------------------------------------------------

/// Transform a pipeline, expanding array literals and handling hoisting.
fn transform_pipeline(in_pipeline: &TypedPipeline) -> Result<Pipeline, Arc<TranslationError>> {
    let valid = in_pipeline.valid_ref()?;
    let mut builder = pipeline().at(in_pipeline.ast.span.clone());
    let mut name_gen = UniqueNameGenerator::new("__expand");

    for cmd in valid.commands.iter() {
        for expanded_cmd in transform_command(cmd, &mut name_gen) {
            builder = builder.command(expanded_cmd);
        }
    }

    Ok(builder.build())
}

/// Transform a single command, returning potentially multiple commands.
///
/// Returns: hoisted LETs + transformed command + cleanup DROPs
fn transform_command(
    cmd: &Arc<TypedCommand>,
    name_gen: &mut UniqueNameGenerator,
) -> Vec<Arc<Command>> {
    let mut alg = ArrayExpansionAlgebra {
        input_schema: cmd.input_schema.clone(),
        name_gen,
        before_commands: Vec::new(),
        after_commands: Vec::new(),
    };

    // Transform expressions in this command using cata_expressions
    let transformed_cmd_ast = cmd.cata_expressions(&mut alg);

    alg.before_commands
        .into_iter()
        .map(Arc::new)
        .chain(std::iter::once(transformed_cmd_ast))
        .chain(alg.after_commands.into_iter().map(Arc::new))
        .collect()
}

// ---------------------------------------------------------------------------
// Expression Transformation Algebra
// ---------------------------------------------------------------------------

/// Algebra for expanding array literals with struct elements.
///
/// Collects before/after commands for hoisting complex expressions.
struct ArrayExpansionAlgebra<'a> {
    input_schema: Arc<TypeEnvironment>,
    name_gen: &'a mut UniqueNameGenerator,
    before_commands: Vec<Command>,
    after_commands: Vec<Command>,
}

impl MapExpressionAlgebra for ArrayExpansionAlgebra<'_> {
    fn array_literal(
        &mut self,
        node: &TypedArrayLiteral,
        expr: &TypedExpression,
        children: Vec<Arc<Expression>>,
    ) -> Arc<Expression> {
        // Check if this array has struct elements that need expansion
        let Type::Array(arr_type) = expr.resolved_type.as_ref() else {
            // Not an array type - just rebuild with transformed children
            return Arc::new(Expression {
                span: expr.ast.span.clone(),
                kind: hamelin_lib::tree::ast::expression::ArrayLiteral { elements: children }
                    .into(),
            });
        };
        let Type::Struct(element_type) = arr_type.element_type.as_ref() else {
            // Not struct elements - just rebuild with transformed children
            return Arc::new(Expression {
                span: expr.ast.span.clone(),
                kind: hamelin_lib::tree::ast::expression::ArrayLiteral { elements: children }
                    .into(),
            });
        };

        // Check if any element needs expansion
        let needs_expansion = node.elements.iter().any(
            |elem| matches!(elem.resolved_type.as_ref(), Type::Struct(s) if s != element_type),
        );

        if !needs_expansion {
            // No expansion needed - just rebuild with transformed children
            return Arc::new(Expression {
                span: expr.ast.span.clone(),
                kind: hamelin_lib::tree::ast::expression::ArrayLiteral { elements: children }
                    .into(),
            });
        }

        // Expand elements that need it, passing transformed children so nested
        // rewrites (e.g., inner array expansions) are preserved for elements
        // that don't need widening.
        let (expanded_ast, before, after) = expand_array_literal(
            node,
            &children,
            element_type,
            &self.input_schema,
            self.name_gen,
        );
        self.before_commands.extend(before);
        self.after_commands.extend(after);

        // Return the expanded AST
        expanded_ast
    }
}

/// Expand array literal elements to match the array's element type.
///
/// Returns `(expanded_ast, before_commands, after_commands)` where:
/// - `expanded_ast` is the transformed AST (untyped - caller must re-typecheck)
/// - `before_commands` are LET commands to insert before the command using this array
/// - `after_commands` are DROP commands to insert after to clean up hoisted variables
fn expand_array_literal(
    arr: &TypedArrayLiteral,
    children: &[Arc<Expression>],
    target_struct: &Struct,
    input_schema: &Arc<TypeEnvironment>,
    name_gen: &mut UniqueNameGenerator,
) -> (Arc<Expression>, Vec<Command>, Vec<Command>) {
    let mut array_builder = array();
    let mut all_before = Vec::new();
    let mut all_after = Vec::new();

    for (elem, child_ast) in arr.elements.iter().zip(children.iter()) {
        if let Type::Struct(elem_struct) = elem.resolved_type.as_ref() {
            if elem_struct != target_struct {
                // Need to expand this element
                let (expanded_ast, before, after) = expand_struct_to_type_with_ast(
                    elem,
                    Some(child_ast),
                    elem_struct,
                    target_struct,
                    name_gen,
                    input_schema,
                );
                all_before.extend(before);
                all_after.extend(after);
                array_builder = array_builder.element(expanded_ast);
            } else {
                // Element already matches - use transformed child to preserve nested rewrites
                array_builder = array_builder.element(child_ast.clone());
            }
        } else {
            // Non-struct element - use transformed child to preserve nested rewrites
            array_builder = array_builder.element(child_ast.clone());
        }
    }

    // Build the new array AST (untyped - will be re-typechecked at pipeline level)
    let new_ast: Expression = array_builder.build();

    (Arc::new(new_ast), all_before, all_after)
}

#[cfg(test)]
mod tests {
    use super::*;
    use hamelin_lib::{
        err::Context,
        func::registry::FunctionRegistry,
        provider::EnvironmentProvider,
        tree::{
            ast::identifier::{Identifier, SimpleIdentifier as AstSimpleIdentifier},
            builder::{
                array, cast, field_ref, let_command, pipeline, query, struct_literal,
                NullLiteralBuilder, QueryBuilderWithMain,
            },
            typed_ast::command::TypedCommandKind,
        },
        type_check_with_provider,
        types::{array::Array, struct_type::Struct, Type, INT, STRING},
    };
    use pretty_assertions::assert_eq;
    use rstest::rstest;

    #[derive(Debug)]
    struct MockProvider;

    impl EnvironmentProvider for MockProvider {
        fn reflect_columns(&self, name: &Identifier) -> anyhow::Result<Struct> {
            let events: Identifier = AstSimpleIdentifier::new("events").into();

            if name == &events {
                Ok(Struct::default()
                    .with_str("x", INT)
                    .with_str("s", Struct::default().with_str("a", INT).into()))
            } else {
                anyhow::bail!("Table not found: {}", name)
            }
        }

        fn reflect_datasets(&self) -> anyhow::Result<Vec<Identifier>> {
            Ok(vec![])
        }
    }

    fn typed_pipeline(builder: QueryBuilderWithMain) -> Arc<TypedPipeline> {
        let statement = type_check_with_provider(builder.build(), Arc::new(MockProvider)).output;
        statement.pipeline.clone()
    }

    fn test_error(message: impl Into<String>) -> Arc<TranslationError> {
        Arc::new(TranslationError::new(Context::new(0..=0, &message.into())))
    }

    fn array_element_types(
        pipeline: &TypedPipeline,
        field_name: &str,
    ) -> Result<Vec<Type>, Arc<TranslationError>> {
        let valid = pipeline.valid_ref()?;
        for cmd in &valid.commands {
            let TypedCommandKind::Let(let_cmd) = &cmd.kind else {
                continue;
            };

            for assignment in &let_cmd.projections.assignments {
                let Ok(identifier) = assignment.identifier.valid_ref() else {
                    continue;
                };
                let Identifier::Simple(simple) = identifier else {
                    continue;
                };
                if simple.as_str() != field_name {
                    continue;
                }

                let TypedExpressionKind::ArrayLiteral(arr) = &assignment.expression.kind else {
                    return Err(test_error("expected array literal assignment"));
                };

                return Ok(arr
                    .elements
                    .iter()
                    .map(|elem| elem.resolved_type.as_ref().clone())
                    .collect());
            }
        }

        Err(test_error("array literal assignment not found"))
    }

    #[test]
    fn test_debug_expansion() -> Result<(), Arc<TranslationError>> {
        let input = pipeline()
            .from(|f| f.table_reference("events"))
            .command(
                let_command()
                    .named_field(
                        "arr",
                        array()
                            .element(struct_literal().field("a", 1))
                            .element(struct_literal().field("a", 2).field("b", 3)),
                    )
                    .build(),
            )
            .build();

        let input_typed = typed_pipeline(query().main(input));

        // Print element types
        let element_types = array_element_types(&input_typed, "arr")?;
        println!("Element types: {:?}", element_types);

        // Get the array's resolved type
        let valid = input_typed.valid_ref()?;
        for cmd in &valid.commands {
            let TypedCommandKind::Let(let_cmd) = &cmd.kind else {
                continue;
            };
            for assignment in &let_cmd.projections.assignments {
                if let TypedExpressionKind::ArrayLiteral(arr) = &assignment.expression.kind {
                    let Type::Array(arr_type) = assignment.expression.resolved_type.as_ref() else {
                        continue;
                    };
                    println!("Array element type: {:?}", arr_type.element_type);

                    // Check each element
                    for (i, elem) in arr.elements.iter().enumerate() {
                        println!(
                            "Element {} type: {:?}, needs expansion: {}",
                            i,
                            elem.resolved_type,
                            elem.resolved_type.as_ref() != arr_type.element_type.as_ref()
                        );
                    }
                }
            }
        }

        // Now run the expansion
        let registry = Arc::new(FunctionRegistry::default());
        let provider = Arc::new(MockProvider);
        let mut ctx = StatementTranslationContext::new(registry, provider);
        let result = expand_array_literals(input_typed, &mut ctx)?;

        println!("Result AST: {:?}", result.ast);

        Ok(())
    }

    #[test]
    fn test_array_literal_element_types() -> Result<(), Arc<TranslationError>> {
        let input = pipeline()
            .from(|f| f.table_reference("events"))
            .command(
                let_command()
                    .named_field(
                        "arr",
                        array()
                            .element(struct_literal().field("a", 1))
                            .element(struct_literal().field("a", 2).field("b", 3)),
                    )
                    .build(),
            )
            .build();

        let input_typed = typed_pipeline(query().main(input));
        let element_types = array_element_types(&input_typed, "arr")?;

        let expected_first = Struct::default().with_str("a", INT).into();
        let expected_second = Struct::default()
            .with_str("a", INT)
            .with_str("b", INT)
            .into();

        assert_eq!(element_types, vec![expected_first, expected_second]);
        Ok(())
    }

    #[rstest]
    // Case 1: No array literals - pass through unchanged
    #[case::no_expansion_needed(
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("y", 1))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("y", 1))
            .build(),
        Struct::default()
            .with_str("y", INT)
            .with_str("x", INT)
            .with_str("s", Struct::default().with_str("a", INT).into())
    )]
    // Case 2: Widen struct element with missing field (struct literal - expanded inline)
    #[case::struct_element_missing_field(
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(struct_literal().field("a", 1))
                        .element(struct_literal().field("a", 2).field("b", 3)),
                )
                .build())
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(struct_literal()
                            .field("a", 1)
                            .field("b", cast(NullLiteralBuilder::new(), INT)))
                        .element(struct_literal().field("a", 2).field("b", 3)),
                )
                .build())
            .build(),
        Struct::default()
            .with_str(
                "arr",
                Array::new(
                    Struct::default().with_str("a", INT).with_str("b", INT).into(),
                )
                .into(),
            )
            .with_str("x", INT)
            .with_str("s", Struct::default().with_str("a", INT).into())
    )]
    // Case 3: Nested struct widening inside array elements (struct literals - expanded inline)
    #[case::nested_struct_widening(
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(
                            struct_literal().field(
                                "nested",
                                struct_literal().field("x", 1),
                            ),
                        )
                        .element(
                            struct_literal().field(
                                "nested",
                                struct_literal().field("x", 2).field("y", 3),
                            ),
                        ),
                )
                .build())
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(
                            struct_literal().field(
                                "nested",
                                struct_literal()
                                    .field("x", 1)
                                    .field("y", cast(NullLiteralBuilder::new(), INT)),
                            ),
                        )
                        .element(
                            struct_literal().field(
                                "nested",
                                struct_literal().field("x", 2).field("y", 3),
                            ),
                        ),
                )
                .build())
            .build(),
        Struct::default()
            .with_str(
                "arr",
                Array::new(
                    Struct::default()
                        .with_str(
                            "nested",
                            Struct::default().with_str("x", INT).with_str("y", INT).into(),
                        )
                        .into(),
                )
                .into(),
            )
            .with_str("x", INT)
            .with_str("s", Struct::default().with_str("a", INT).into())
    )]
    // Case 4: Simple cast of struct literal (no hoisting since all fields are simple)
    #[case::simple_cast_struct_literal(
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(cast(
                            struct_literal().field("a", 1).field("b", 2),
                            Struct::default().with_str("a", INT).with_str("b", INT).into(),
                        ))
                        .element(
                            struct_literal()
                                .field("a", 3)
                                .field("b", 4)
                                .field("c", 5),
                        ),
                )
                .build())
            .build(),
        // No hoisting because cast(struct_literal()) is simple
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(
                            // Simple expr is just cast to target type (no hoisting)
                            cast(
                                cast(
                                    struct_literal().field("a", 1).field("b", 2),
                                    Struct::default().with_str("a", INT).with_str("b", INT).into(),
                                ),
                                Struct::default()
                                    .with_str("a", INT)
                                    .with_str("b", INT)
                                    .with_str("c", INT)
                                    .into(),
                            ),
                        )
                        .element(
                            struct_literal()
                                .field("a", 3)
                                .field("b", 4)
                                .field("c", 5),
                        ),
                )
                .build())
            .build(),
        Struct::default()
            .with_str(
                "arr",
                Array::new(
                    Struct::default()
                        .with_str("a", INT)
                        .with_str("b", INT)
                        .with_str("c", INT)
                        .into(),
                )
                .into(),
            )
            .with_str("x", INT)
            .with_str("s", Struct::default().with_str("a", INT).into())
    )]
    // Case 5: Column reference chain expansion (emits cast, no hoisting)
    #[case::field_reference_chain_expansion(
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(field_ref("s"))  // s: {a: Int} from table
                        .element(struct_literal().field("a", 1).field("b", 2)),
                )
                .build())
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(
                            // Column ref is cast to target type
                            cast(
                                field_ref("s"),
                                Struct::default().with_str("a", INT).with_str("b", INT).into(),
                            ),
                        )
                        .element(struct_literal().field("a", 1).field("b", 2)),
                )
                .build())
            .build(),
        Struct::default()
            .with_str(
                "arr",
                Array::new(
                    Struct::default().with_str("a", INT).with_str("b", INT).into(),
                )
                .into(),
            )
            .with_str("x", INT)
            .with_str("s", Struct::default().with_str("a", INT).into())
    )]
    // Case 6: Simple cast expression (emits cast, no hoisting since cast is simple)
    #[case::simple_cast_no_hoisting(
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(cast(
                            struct_literal().field("a", 1),
                            Struct::default().with_str("a", INT).into(),
                        ))
                        .element(struct_literal().field("a", 2).field("b", 3)),
                )
                .build())
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(
                            // Cast expression is simple enough, just wrap in another cast
                            cast(
                                cast(
                                    struct_literal().field("a", 1),
                                    Struct::default().with_str("a", INT).into(),
                                ),
                                Struct::default().with_str("a", INT).with_str("b", INT).into(),
                            ),
                        )
                        .element(struct_literal().field("a", 2).field("b", 3)),
                )
                .build())
            .build(),
        Struct::default()
            .with_str(
                "arr",
                Array::new(
                    Struct::default().with_str("a", INT).with_str("b", INT).into(),
                )
                .into(),
            )
            .with_str("x", INT)
            .with_str("s", Struct::default().with_str("a", INT).into())
    )]
    // Case 7: Array with explicit null AS string casts (no expansion needed - all same type)
    #[case::explicit_null_casts_preserved(
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(struct_literal().field("left", "apple").field("right", "zebra"))
                        .element(struct_literal().field("left", cast(NullLiteralBuilder::new(), STRING)).field("right", "yak"))
                        .element(struct_literal().field("left", "cherry").field("right", cast(NullLiteralBuilder::new(), STRING)))
                        .element(struct_literal().field("left", "date").field("right", "walrus")),
                )
                .build())
            .build(),
        // Expected: unchanged - all elements already have the same struct type
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(struct_literal().field("left", "apple").field("right", "zebra"))
                        .element(struct_literal().field("left", cast(NullLiteralBuilder::new(), STRING)).field("right", "yak"))
                        .element(struct_literal().field("left", "cherry").field("right", cast(NullLiteralBuilder::new(), STRING)))
                        .element(struct_literal().field("left", "date").field("right", "walrus")),
                )
                .build())
            .build(),
        Struct::default()
            .with_str(
                "arr",
                Array::new(
                    Struct::default().with_str("left", STRING).with_str("right", STRING).into(),
                )
                .into(),
            )
            .with_str("x", INT)
            .with_str("s", Struct::default().with_str("a", INT).into())
    )]
    // Case 8: Nested array-in-struct widening (emits cast for array field)
    #[case::nested_array_in_struct_widening(
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        // First element: {items: [{a: 1}]}
                        .element(struct_literal().field(
                            "items",
                            array().element(struct_literal().field("a", 1)),
                        ))
                        // Second element: {items: [{a: 2, b: 3}]}
                        .element(struct_literal().field(
                            "items",
                            array().element(struct_literal().field("a", 2).field("b", 3)),
                        )),
                )
                .build())
            .build(),
        // Expected: first element's items field is cast to array<{a, b}>
        pipeline()
            .from(|f| f.table_reference("events"))
            .command(let_command()
                .named_field(
                    "arr",
                    array()
                        .element(struct_literal().field(
                            "items",
                            cast(
                                array().element(struct_literal().field("a", 1)),
                                Array::new(
                                    Struct::default().with_str("a", INT).with_str("b", INT).into(),
                                )
                                .into(),
                            ),
                        ))
                        .element(struct_literal().field(
                            "items",
                            array().element(struct_literal().field("a", 2).field("b", 3)),
                        )),
                )
                .build())
            .build(),
        Struct::default()
            .with_str(
                "arr",
                Array::new(
                    Struct::default()
                        .with_str(
                            "items",
                            Array::new(
                                Struct::default().with_str("a", INT).with_str("b", INT).into(),
                            )
                            .into(),
                        )
                        .into(),
                )
                .into(),
            )
            .with_str("x", INT)
            .with_str("s", Struct::default().with_str("a", INT).into())
    )]
    fn test_expand_array_literals(
        #[case] input: Pipeline,
        #[case] expected: Pipeline,
        #[case] expected_output_schema: Struct,
    ) -> Result<(), Arc<TranslationError>> {
        let input_typed = typed_pipeline(query().main(input));
        let expected_typed = typed_pipeline(query().main(expected));

        let registry = Arc::new(FunctionRegistry::default());
        let provider = Arc::new(MockProvider);
        let mut ctx = StatementTranslationContext::new(registry, provider);
        let result = expand_array_literals(input_typed, &mut ctx)?;

        assert_eq!(result.ast, expected_typed.ast);
        assert_eq!(
            result.environment().as_struct().clone(),
            expected_output_schema
        );
        Ok(())
    }
}