hamelin_translation 0.9.6

//! Pipeline pass: Extract nested aggregates from WINDOW expressions.
//!
//! Transforms WINDOW commands where aggregate functions are buried in expressions
//! into WINDOW commands with only top-level aggregates, followed by SET + DROP
//! that compute the final expression.
//!
//! Example:
//! ```text
//! WINDOW crazy_number = sum(left) + sum(right) + 3 SORT order
//! ```
//! becomes:
//! ```text
//! WINDOW __window_agg_0 = sum(left), __window_agg_1 = sum(right) SORT order
//! | SET crazy_number = __window_agg_0 + __window_agg_1 + 3
//! | DROP __window_agg_0, __window_agg_1
//! ```
//!
//! This ensures the IR's WINDOW command only contains top-level aggregate functions,
//! making backend translation straightforward.

use std::sync::Arc;

use hamelin_lib::err::TranslationError;
use hamelin_lib::tree::ast::expression::Expression;
use hamelin_lib::tree::ast::identifier::{Identifier, SimpleIdentifier};
use hamelin_lib::tree::typed_ast::command::{TypedCommand, TypedCommandKind, TypedWindowCommand};
use hamelin_lib::tree::typed_ast::context::StatementTranslationContext;
use hamelin_lib::tree::typed_ast::pipeline::TypedPipeline;
use hamelin_lib::tree::{
    ast::command::Command,
    builder::{self, drop_command, set_command, window_command},
};

use crate::normalize::special_function_extraction::{
    has_nested_special_function, ExtractSpecialFunctionsAlgebra,
};
use crate::unique::UniqueNameGenerator;

/// Extract nested aggregates from WINDOW expressions.
///
/// Contract: `Arc<TypedPipeline> -> Result<Arc<TypedPipeline>, ...>`
pub fn extract_window_aggregates(
    pipeline: Arc<TypedPipeline>,
    ctx: &mut StatementTranslationContext,
) -> Result<Arc<TypedPipeline>, Arc<TranslationError>> {
    // Check if any WINDOW command has nested aggregates
    if !pipeline
        .valid_ref()?
        .commands
        .iter()
        .any(window_has_nested_aggregates)
    {
        return Ok(pipeline);
    }

    let valid = pipeline.valid_ref()?;

    // Shared name generator
    let mut name_gen = UniqueNameGenerator::new("__window_agg");

    // Transform commands
    let mut pipe_builder = builder::pipeline();
    for cmd in &valid.commands {
        for c in transform_command(cmd, &mut name_gen) {
            pipe_builder = pipe_builder.command(c);
        }
    }

    let new_ast = pipe_builder.build().at(pipeline.ast.span);

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

/// Check if a command is a WINDOW with nested aggregates in projections.
fn window_has_nested_aggregates(cmd: &Arc<TypedCommand>) -> bool {
    let TypedCommandKind::Window(window_cmd) = &cmd.kind else {
        return false;
    };
    window_cmd
        .projections
        .assignments
        .iter()
        .any(|a| has_nested_special_function(&a.expression))
}

/// Transform a single command.
fn transform_command(
    cmd: &Arc<TypedCommand>,
    name_gen: &mut UniqueNameGenerator,
) -> Vec<Arc<Command>> {
    let TypedCommandKind::Window(window_cmd) = &cmd.kind else {
        return vec![cmd.ast.clone()];
    };

    if !window_has_nested_aggregates(cmd) {
        return vec![cmd.ast.clone()];
    }

    transform_window(window_cmd, cmd, name_gen)
}

/// Transform a WINDOW command, extracting nested aggregates.
fn transform_window(
    window_cmd: &TypedWindowCommand,
    cmd: &TypedCommand,
    name_gen: &mut UniqueNameGenerator,
) -> Vec<Arc<Command>> {
    let mut extracted_aggs: Vec<(SimpleIdentifier, Expression)> = Vec::new();
    let mut set_assignments: Vec<(SimpleIdentifier, Expression)> = Vec::new();
    let mut synth_ids_to_drop: Vec<SimpleIdentifier> = Vec::new();

    // Process each projection
    for assignment in &window_cmd.projections.assignments {
        let Ok(id) = assignment.identifier.valid_ref() else {
            continue;
        };
        let simple_id = match id {
            Identifier::Simple(s) => s.clone(),
            Identifier::Compound(c) => {
                // Should have been normalized by normalize_window already
                SimpleIdentifier::new(c.to_string())
            }
        };

        if has_nested_special_function(&assignment.expression) {
            // Extract special-position calls using cata and build replacement expression
            let mut alg = ExtractSpecialFunctionsAlgebra {
                name_gen,
                schema: &cmd.input_schema,
                extractions: Vec::new(),
                synth_ids: Vec::new(),
            };
            let new_ast = assignment.expression.cata(&mut alg);

            extracted_aggs.extend(alg.extractions);
            synth_ids_to_drop.extend(alg.synth_ids);
            set_assignments.push((simple_id, new_ast.as_ref().clone()));
        } else {
            // Top-level aggregate or no aggregate - keep in WINDOW
            extracted_aggs.push((simple_id, assignment.expression.ast.as_ref().clone()));
        }
    }

    // Build the WINDOW command with extracted aggregates
    let mut window_builder = window_command().at(cmd.ast.span);
    for (id, expr) in &extracted_aggs {
        window_builder = window_builder.named_field(id.clone(), expr.clone());
    }

    // Preserve group_by
    for assignment in &window_cmd.group_by.assignments {
        if let Ok(id) = assignment.identifier.valid_ref() {
            window_builder =
                window_builder.group_by(id.clone(), assignment.expression.ast.as_ref().clone());
        }
    }

    // Preserve sort_by
    for sort_expr in &window_cmd.sort_by {
        window_builder = window_builder.sort_expr(sort_expr.ast.as_ref().clone());
    }

    // Preserve within
    if let Some(within) = &window_cmd.within {
        window_builder = window_builder.within(within.ast.clone());
    }

    let mut result = vec![Arc::new(window_builder.build())];

    // If we extracted any aggregates into expressions, add SET + DROP commands
    if !set_assignments.is_empty() {
        // SET to compute the final expressions from the synthetic aggregates
        let mut set_builder = set_command();
        for (id, expr) in &set_assignments {
            set_builder = set_builder.named_field(id.clone(), expr.clone());
        }
        result.push(Arc::new(set_builder.build()));

        // DROP to remove the synthetic aggregate columns
        let mut drop_builder = drop_command();
        for synth_id in &synth_ids_to_drop {
            drop_builder = drop_builder.field(synth_id.clone());
        }
        result.push(Arc::new(drop_builder.build()));
    }

    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use hamelin_lib::type_check;
    use hamelin_lib::{
        tree::{
            ast::pipeline::Pipeline,
            builder::{
                add, call, drop_command, field_ref, pipeline, select_command, set_command,
                sort_command, window_command,
            },
        },
        types::{struct_type::Struct, INT},
    };
    use pretty_assertions::assert_eq;
    use rstest::rstest;
    use std::sync::Arc;

    #[rstest]
    // Case 1: No WINDOW commands - passes through unchanged
    #[case::no_window_passthrough(
        pipeline()
            .command(select_command().named_field("a", 1).named_field("b", 2).build())
            .build(),
        pipeline()
            .command(select_command().named_field("a", 1).named_field("b", 2).build())
            .build(),
        Struct::default().with_str("a", INT).with_str("b", INT)
    )]
    // Case 2: WINDOW with top-level aggregate - passes through unchanged
    #[case::window_toplevel_agg_unchanged(
        pipeline()
            .command(select_command().named_field("value", 10).named_field("order", 1).build())
            .command(window_command()
                .named_field("total", call("sum").arg(field_ref("value")))
                .sort(sort_command().by(field_ref("order")))
                .build())
            .build(),
        pipeline()
            .command(select_command().named_field("value", 10).named_field("order", 1).build())
            .command(window_command()
                .named_field("total", call("sum").arg(field_ref("value")))
                .sort(sort_command().by(field_ref("order")))
                .build())
            .build(),
        Struct::default()
            .with_str("total", INT)
            .with_str("value", INT)
            .with_str("order", INT)
    )]
    // Case 3: WINDOW with nested aggregates in binary expression
    // Input: WINDOW crazy_number = sum(left) + sum(right) + 3
    // Output: WINDOW __window_agg_0 = sum(left), __window_agg_1 = sum(right)
    //       | SET crazy_number = __window_agg_0 + __window_agg_1 + 3
    //       | DROP __window_agg_0, __window_agg_1
    #[case::window_nested_agg_binary(
        pipeline()
            .command(select_command()
                .named_field("left", 1)
                .named_field("right", 2)
                .named_field("order", 1)
                .build())
            .command(window_command()
                .named_field("crazy_number",
                    add(
                        add(
                            call("sum").arg(field_ref("left")),
                            call("sum").arg(field_ref("right"))
                        ),
                        3
                    )
                )
                .sort(sort_command().by(field_ref("order")))
                .build())
            .build(),
        pipeline()
            .command(select_command()
                .named_field("left", 1)
                .named_field("right", 2)
                .named_field("order", 1)
                .build())
            .command(window_command()
                .named_field("__window_agg_0", call("sum").arg(field_ref("left")))
                .named_field("__window_agg_1", call("sum").arg(field_ref("right")))
                .sort(sort_command().by(field_ref("order")))
                .build())
            .command(set_command()
                .named_field("crazy_number",
                    add(
                        add(
                            field_ref("__window_agg_0"),
                            field_ref("__window_agg_1")
                        ),
                        3
                    )
                )
                .build())
            .command(drop_command()
                .field("__window_agg_0")
                .field("__window_agg_1")
                .build())
            .build(),
        Struct::default()
            .with_str("crazy_number", INT)
            .with_str("left", INT)
            .with_str("right", INT)
            .with_str("order", INT)
    )]
    fn test_extract_window_aggregates(
        #[case] input: Pipeline,
        #[case] expected: Pipeline,
        #[case] expected_output_schema: Struct,
    ) {
        let input_typed = type_check(input).output;
        let expected_typed = type_check(expected).output;

        let mut ctx = StatementTranslationContext::default();
        let result = extract_window_aggregates(Arc::new(input_typed), &mut ctx).unwrap();

        // Compare ASTs
        assert_eq!(result.ast, expected_typed.ast);

        // Verify output schema
        let result_schema = result.environment().as_struct().clone();
        assert_eq!(result_schema, expected_output_schema);
    }
}