hamelin_translation 0.4.2

//! Pipeline pass: Broadcast lowering.
//!
//! Transforms broadcast operations into explicit `transform(array, lambda)` calls.
//!
//! This handles two kinds of broadcasts:
//!
//! 1. **BroadcastApply**: A function call where one argument is an array and the
//!    operation should be applied element-wise.
//!    ```text
//!    [1, 2, 3] * 10  →  transform([1, 2, 3], __broadcast_0 -> __broadcast_0 * 10)
//!    ```
//!
//! 2. **FieldLookup on Array**: Field access on an array broadcasts over elements.
//!    ```text
//!    events.user  →  transform(events, __broadcast_0 -> __broadcast_0.user)
//!    ```
//!    where `events` is `Array<Struct>` or `Array<Variant>`.

use std::sync::Arc;

use hamelin_lib::{
    err::TranslationError,
    tree::{
        ast::expression::Expression,
        builder::{call, column_ref, field, lambda1, pipeline, ExpressionBuilder},
        typed_ast::{
            command::TypedCommand,
            context::StatementTranslationContext,
            environment::TypeEnvironment,
            expression::{
                MapExpressionAlgebra, TypedBroadcastApply, TypedExpression, TypedFieldLookup,
            },
            pipeline::TypedPipeline,
        },
    },
    types::Type,
};

use super::super::unique::UniqueNameGenerator;

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

/// Lower BroadcastApply nodes to explicit transform() calls.
///
/// Contract: `(Arc<TypedPipeline>, &mut ctx) -> Result<Arc<TypedPipeline>, ...>`
pub fn lower_broadcast_apply(
    pipeline: Arc<TypedPipeline>,
    ctx: &mut StatementTranslationContext,
) -> Result<Arc<TypedPipeline>, Arc<TranslationError>> {
    // Check if any command has BroadcastApply
    if !pipeline_has_broadcast(&pipeline)? {
        return Ok(pipeline);
    }

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

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

// ---------------------------------------------------------------------------
// Detection
// ---------------------------------------------------------------------------

/// Check if a pipeline has any BroadcastApply nodes.
fn pipeline_has_broadcast(pipeline: &TypedPipeline) -> Result<bool, Arc<TranslationError>> {
    let valid = pipeline.valid_ref()?;
    Ok(valid.commands.iter().any(command_has_broadcast))
}

/// Check if a command has any broadcast operations (BroadcastApply or FieldLookup on array).
fn command_has_broadcast(cmd: &Arc<TypedCommand>) -> bool {
    use hamelin_lib::tree::typed_ast::expression::TypedExpressionKind;

    cmd.find_expression(&mut |expr| match &expr.kind {
        TypedExpressionKind::BroadcastApply(_) => true,
        TypedExpressionKind::FieldLookup(fl) => {
            matches!(fl.value.resolved_type.as_ref(), Type::Array(_))
        }
        _ => false,
    })
    .is_some()
}

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

/// Transform a pipeline, lowering all BroadcastApply nodes.
fn transform_pipeline(
    in_pipeline: &TypedPipeline,
) -> Result<hamelin_lib::tree::ast::pipeline::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("__broadcast");

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

    Ok(builder.build())
}

/// Transform a single command, lowering BroadcastApply nodes within.
fn transform_command(
    cmd: &Arc<TypedCommand>,
    name_gen: &mut UniqueNameGenerator,
) -> Arc<hamelin_lib::tree::ast::command::Command> {
    let mut alg = BroadcastLoweringAlgebra {
        name_gen,
        schema: &cmd.input_schema,
    };
    cmd.cata_expressions(&mut alg)
}

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

/// Algebra for lowering broadcast operations to transform() calls.
struct BroadcastLoweringAlgebra<'a> {
    name_gen: &'a mut UniqueNameGenerator,
    schema: &'a TypeEnvironment,
}

impl MapExpressionAlgebra for BroadcastLoweringAlgebra<'_> {
    fn broadcast_apply(
        &mut self,
        node: &TypedBroadcastApply,
        expr: &TypedExpression,
        children: hamelin_lib::func::def::ParameterBinding<Arc<Expression>>,
    ) -> Arc<Expression> {
        // Generate a unique parameter name for the lambda
        let param_name = self.name_gen.next(self.schema);

        // Get the array argument by its index position
        let array_arg = children
            .get_by_index(node.broadcast_position)
            .expect("broadcast_position should be valid")
            .clone();

        // Build the lambda body by replacing the broadcast argument with a column
        // reference to the lambda parameter, preserving the original AST structure
        let param_ref = Arc::new(column_ref(param_name.clone()).build());
        let body_children = children
            .replace_by_index(node.broadcast_position, param_ref)
            .expect("broadcast_position should be valid");

        // Use the existing replace_children_ast to rebuild with original AST structure
        // (preserves BinaryOperator vs FunctionCall distinction)
        let body_ast = node.replace_children_ast(expr, body_children);

        // Build the lambda: param -> body
        let lambda_ast = lambda1(param_name.as_str())
            .body(AstExpressionWrapper(body_ast))
            .build();

        // Build transform(array, lambda)
        let transform_call = call("transform")
            .arg(AstExpressionWrapper(array_arg))
            .arg(lambda_ast)
            .build();

        Arc::new(transform_call)
    }

    fn field_lookup(
        &mut self,
        node: &TypedFieldLookup,
        expr: &TypedExpression,
        child: Arc<Expression>,
    ) -> Arc<Expression> {
        // Only transform if the value is an array (broadcast case)
        if !matches!(node.value.resolved_type.as_ref(), Type::Array(_)) {
            // Not an array - use default behavior (rebuild AST)
            return node.replace_children_ast(expr, child);
        }

        // Extract the field name from the original AST
        let field_name = match &expr.ast.kind {
            hamelin_lib::tree::ast::expression::ExpressionKind::FieldLookup(fl) => {
                // Get the field name as a string
                fl.field_identifier.to_string()
            }
            _ => return node.replace_children_ast(expr, child),
        };

        // Generate a unique parameter name for the lambda
        let param_name = self.name_gen.next(self.schema);

        // Build the lambda body: param.field
        let body_ast = field(column_ref(param_name.clone()), &field_name).build();

        // Build the lambda: param -> param.field
        let lambda_ast = lambda1(param_name.as_str()).body(body_ast).build();

        // Build transform(array, lambda)
        let transform_call = call("transform")
            .arg(AstExpressionWrapper(child))
            .arg(lambda_ast)
            .build();

        Arc::new(transform_call)
    }
}

/// Wrapper to use an existing AST Expression as an ExpressionBuilder.
#[derive(Debug)]
struct AstExpressionWrapper(Arc<Expression>);

impl ExpressionBuilder for AstExpressionWrapper {
    fn build(&self) -> Expression {
        self.0.as_ref().clone()
    }
}

impl hamelin_lib::tree::builder::IntoExpressionBuilder for AstExpressionWrapper {
    fn into_expression_builder(self) -> Box<dyn ExpressionBuilder> {
        Box::new(self)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use hamelin_lib::{
        provider::EnvironmentProvider,
        sql::{
            expression::identifier::Identifier as SqlIdentifier,
            query::TableReference as SqlTableReference,
        },
        tree::{
            ast::{pipeline::Pipeline, IntoTyped, TypeCheckExecutor},
            builder::{
                add, array, call, column_ref, gt, lambda1, multiply, pipeline, query, HasMain,
                QueryBuilder,
            },
        },
        types::{array::Array, struct_type::Struct, Type, INT, STRING},
    };
    use pretty_assertions::assert_eq;
    use rstest::rstest;
    use std::sync::Arc;

    #[derive(Debug)]
    struct MockProvider;

    impl EnvironmentProvider for MockProvider {
        fn reflect_columns(&self, table: SqlTableReference) -> anyhow::Result<Struct> {
            let events: SqlIdentifier = "events".parse().unwrap();

            if table.name == events {
                Ok(Struct::default()
                    .with_str("x", INT)
                    .with_str("s", STRING)
                    .with_str("arr_col", Array::new(INT).into()))
            } else {
                anyhow::bail!("Table not found: {}", table.name)
            }
        }

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

    fn typed_pipeline(builder: QueryBuilder<HasMain>) -> Arc<TypedPipeline> {
        let statement = builder
            .build()
            .typed_with()
            .with_provider(Arc::new(MockProvider))
            .typed();
        statement.pipeline.clone()
    }

    fn make_ctx() -> StatementTranslationContext {
        StatementTranslationContext::new(
            Arc::new(hamelin_lib::func::registry::FunctionRegistry::default()),
            Arc::new(MockProvider),
        )
    }

    /// Base schema from events table: x, s, arr_col
    fn base_schema() -> Struct {
        Struct::default()
            .with_str("x", INT)
            .with_str("s", STRING)
            .with_str("arr_col", Array::new(INT).into())
    }

    /// Schema with new field prepended (LET adds fields at beginning)
    fn schema_with_let(field: &str, typ: Type) -> Struct {
        Struct::default()
            .with_str(field, typ)
            .with_str("x", INT)
            .with_str("s", STRING)
            .with_str("arr_col", Array::new(INT).into())
    }

    #[rstest]
    // Case 1: No broadcast - passes through unchanged
    #[case::no_broadcast_passthrough(
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("y", add(column_ref("x"), 1)))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("y", add(column_ref("x"), 1)))
            .build(),
        schema_with_let("y", INT)
    )]
    // Case 2: Simple binary broadcast on left: [1,2,3] * 10
    // The lambda body preserves the BinaryOperator AST structure
    #[case::left_broadcast(
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("arr", multiply(array().element(1).element(2).element(3), 10)))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("arr", call("transform")
                .arg(array().element(1).element(2).element(3))
                .arg(lambda1("__broadcast_0").body(multiply(column_ref("__broadcast_0"), 10)))))
            .build(),
        schema_with_let("arr", Array::new(INT).into())
    )]
    // Case 3: Binary broadcast on right: 10 * [1,2,3]
    #[case::right_broadcast(
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("arr", multiply(10, array().element(1).element(2).element(3))))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("arr", call("transform")
                .arg(array().element(1).element(2).element(3))
                .arg(lambda1("__broadcast_0").body(multiply(10, column_ref("__broadcast_0"))))))
            .build(),
        schema_with_let("arr", Array::new(INT).into())
    )]
    // Case 4: Function broadcast: upper(["a","b"])
    #[case::function_broadcast(
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("arr", call("upper").arg(array().element("a").element("b"))))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("arr", call("transform")
                .arg(array().element("a").element("b"))
                .arg(lambda1("__broadcast_0").body(call("upper").arg(column_ref("__broadcast_0"))))))
            .build(),
        schema_with_let("arr", Array::new(STRING).into())
    )]
    // Case 5: Two separate broadcasts in same command
    #[case::two_broadcasts_same_command(
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l
                .named_field("a", add(array().element(1).element(2), 5))
                .named_field("b", multiply(array().element(10).element(20), 2)))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l
                .named_field("a", call("transform")
                    .arg(array().element(1).element(2))
                    .arg(lambda1("__broadcast_0").body(add(column_ref("__broadcast_0"), 5))))
                .named_field("b", call("transform")
                    .arg(array().element(10).element(20))
                    .arg(lambda1("__broadcast_1").body(multiply(column_ref("__broadcast_1"), 2)))))
            .build(),
        // Two fields prepended: a first, then b
        Struct::default()
            .with_str("a", Array::new(INT).into())
            .with_str("b", Array::new(INT).into())
            .with_str("x", INT)
            .with_str("s", STRING)
            .with_str("arr_col", Array::new(INT).into())
    )]
    // Case 6: Broadcast over column reference (not literal array)
    // arr_col * 10 where arr_col is Array<Int> from table
    #[case::broadcast_over_column(
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("result", multiply(column_ref("arr_col"), 10)))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("result", call("transform")
                .arg(column_ref("arr_col"))
                .arg(lambda1("__broadcast_0").body(multiply(column_ref("__broadcast_0"), 10)))))
            .build(),
        schema_with_let("result", Array::new(INT).into())
    )]
    // Case 7: Broadcast nested inside a non-broadcasting outer expression
    // coalesce(upper(["a","b"]), ["default"]) - the upper broadcasts, coalesce does not
    // Note: coalesce with two Array<String> args doesn't broadcast, just returns first non-null
    #[case::broadcast_nested_in_coalesce(
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("arr", call("coalesce")
                .arg(call("upper").arg(array().element("a").element("b")))
                .arg(array().element("default"))))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("arr", call("coalesce")
                .arg(call("transform")
                    .arg(array().element("a").element("b"))
                    .arg(lambda1("__broadcast_0").body(call("upper").arg(column_ref("__broadcast_0")))))
                .arg(array().element("default"))))
            .build(),
        schema_with_let("arr", Array::new(STRING).into())
    )]
    // Case 8: Broadcast in WHERE command (not just LET)
    // WHERE len(arr_col * 10) > 0 - broadcast inside WHERE, result used in comparison
    // The arr_col * 10 broadcasts, len takes the array, then > 0 is a scalar comparison
    #[case::broadcast_in_where(
        pipeline()
            .from(|f| f.table_reference("events"))
            .where_cmd(gt(call("len").arg(multiply(column_ref("arr_col"), 10)), 0))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .where_cmd(gt(
                call("len").arg(
                    call("transform")
                        .arg(column_ref("arr_col"))
                        .arg(lambda1("__broadcast_0").body(multiply(column_ref("__broadcast_0"), 10)))),
                0))
            .build(),
        // WHERE doesn't add fields
        base_schema()
    )]
    // Case 9: Multiple commands with broadcasts - counter increments across commands
    #[case::broadcasts_across_commands(
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("a", multiply(array().element(1).element(2), 5)))
            .let_cmd(|l| l.named_field("b", add(array().element(10).element(20), 3)))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("a", call("transform")
                .arg(array().element(1).element(2))
                .arg(lambda1("__broadcast_0").body(multiply(column_ref("__broadcast_0"), 5)))))
            .let_cmd(|l| l.named_field("b", call("transform")
                .arg(array().element(10).element(20))
                .arg(lambda1("__broadcast_1").body(add(column_ref("__broadcast_1"), 3)))))
            .build(),
        // First LET adds "a", second LET adds "b" at beginning
        Struct::default()
            .with_str("b", Array::new(INT).into())
            .with_str("a", Array::new(INT).into())
            .with_str("x", INT)
            .with_str("s", STRING)
            .with_str("arr_col", Array::new(INT).into())
    )]
    // Case 10: Name collision - input has column named __broadcast_0
    // The generator now skips existing names, so it will use __broadcast_1 instead.
    #[case::name_collision_avoided(
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("__broadcast_0", 42))  // Create column with generated name
            .let_cmd(|l| l.named_field("result", multiply(column_ref("arr_col"), 10)))
            .build(),
        pipeline()
            .from(|f| f.table_reference("events"))
            .let_cmd(|l| l.named_field("__broadcast_0", 42))
            .let_cmd(|l| l.named_field("result", call("transform")
                .arg(column_ref("arr_col"))
                // Generator skips __broadcast_0 (exists) and uses __broadcast_1
                .arg(lambda1("__broadcast_1").body(multiply(column_ref("__broadcast_1"), 10)))))
            .build(),
        // __broadcast_0 added first, then result
        Struct::default()
            .with_str("result", Array::new(INT).into())
            .with_str("__broadcast_0", INT)
            .with_str("x", INT)
            .with_str("s", STRING)
            .with_str("arr_col", Array::new(INT).into())
    )]
    fn test_lower_broadcast_apply(
        #[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 mut ctx = make_ctx();
        let result = lower_broadcast_apply(input_typed, &mut ctx)?;

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