hamelin_translation 0.3.10

//! Fuse projection commands (SELECT, LET, DROP) into minimal SELECT commands.
//!
//! This pass runs after normalize_commands.
//! Pipeline-level pass contract: `Rc<TypedPipeline> -> Result<Rc<TypedPipeline>, ...>`

use std::collections::HashSet;
use std::rc::Rc;

use ordermap::OrderMap;

use hamelin_lib::{
    err::TranslationError,
    tree::{
        ast::{
            command::Command,
            expression::Expression,
            identifier::{Identifier, SimpleIdentifier},
            node::Span,
        },
        builder::{self, column_ref, field, select_command, ExpressionBuilder},
        typed_ast::{
            clause::Projections,
            command::{
                TypedCommand, TypedCommandKind, TypedDropCommand, TypedLetCommand,
                TypedSelectCommand,
            },
            context::StatementTranslationContext,
            expression::{TypedExpression, TypedExpressionKind},
            pipeline::TypedPipeline,
        },
    },
};

/// Fuse projection commands (SELECT, LET, DROP) into minimal SELECT commands.
///
/// Pipeline-level pass contract: `Rc<TypedPipeline> -> Result<Rc<TypedPipeline>, ...>`
///
/// This creates a minimal number of SELECT commands by:
/// 1. Accumulating LET/DROP operations into a pending SELECT
/// 2. Detecting dependency barriers when LET references a field we assigned
/// 3. Emitting the pending SELECT when we hit a barrier or non-projection command
///
/// Returns the original `Rc` unchanged if no fusion is needed.
pub fn fuse_projections(
    pipeline: Rc<TypedPipeline>,
    ctx: &mut StatementTranslationContext,
) -> Result<Rc<TypedPipeline>, Rc<TranslationError>> {
    // Check if fusion is needed
    let needs_fusion = pipeline
        .valid_ref()?
        .commands
        .iter()
        .any(command_needs_fusion);

    if !needs_fusion {
        return Ok(pipeline);
    }

    let valid = pipeline.valid_ref()?;

    // Fuse commands
    let fused_commands = fuse_commands(&valid.commands)?;

    let mut pipe_builder = builder::pipeline().at(pipeline.ast.span.clone());
    for cmd in fused_commands {
        pipe_builder = pipe_builder.command(cmd);
    }
    let fused_ast = pipe_builder.build();

    // Re-typecheck
    Ok(Rc::new(TypedPipeline::from_ast_with_context(
        Rc::new(fused_ast),
        ctx,
    )))
}

/// Check if a command needs fusion (LET or DROP).
fn command_needs_fusion(cmd: &Rc<TypedCommand>) -> bool {
    matches!(
        &cmd.kind,
        TypedCommandKind::Let(_) | TypedCommandKind::Drop(_)
    )
}

/// Fuse commands into minimal SELECT commands.
fn fuse_commands(commands: &[Rc<TypedCommand>]) -> Result<Vec<Command>, Rc<TranslationError>> {
    let mut result: Vec<Command> = Vec::new();
    let mut pending: Option<PendingSelect> = None;

    for command in commands {
        match &command.kind {
            TypedCommandKind::Select(select_cmd) => {
                // SELECT is a barrier - emit pending and start new
                if let Some(p) = pending.take() {
                    result.push(p.emit());
                }
                pending = Some(PendingSelect::from_select(command, select_cmd)?);
            }

            TypedCommandKind::Let(let_cmd) => {
                let refs = extract_column_references_from_projections(&let_cmd.projections);

                if let Some(ref mut p) = pending {
                    // Check for dependency barrier
                    if refs
                        .iter()
                        .any(|r| p.assigned.iter().any(|a| identifiers_overlap(r, a)))
                    {
                        // BARRIER: LET references something we assigned
                        result.push(pending.take().unwrap().emit());
                        pending = Some(PendingSelect::from_let(command, let_cmd)?);
                    } else {
                        // FUSE: merge LET into pending SELECT
                        p.merge_let(command, let_cmd)?;
                    }
                } else {
                    pending = Some(PendingSelect::from_let(command, let_cmd)?);
                }
            }

            TypedCommandKind::Drop(drop_cmd) => {
                if let Some(ref mut p) = pending {
                    // DROP can always fuse - it just removes fields
                    p.merge_drop(command, drop_cmd);
                } else {
                    pending = Some(PendingSelect::from_drop(command, drop_cmd));
                }
            }

            _ => {
                // Non-projection command (WHERE, SORT, AGG, etc.) is a barrier
                if let Some(p) = pending.take() {
                    result.push(p.emit());
                }
                // Pass through the original AST command
                result.push(command.ast.as_ref().clone());
            }
        }
    }

    // Flush remaining
    if let Some(p) = pending {
        result.push(p.emit());
    }

    Ok(result)
}

/// Pending SELECT being accumulated during fusion.
///
/// Tracks assignments by identifier, preserving insertion order via OrderMap.
struct PendingSelect {
    /// The SELECT assignments being built (identifier -> AST expression)
    /// OrderMap preserves insertion order for deterministic output
    assignments: OrderMap<Identifier, Rc<Expression>>,
    /// Set of identifiers that have been explicitly assigned (not passthrough)
    assigned: HashSet<Identifier>,
    /// AST span for the synthesized command
    span: Span,
}

impl PendingSelect {
    /// Initialize from a SELECT command
    fn from_select(
        command: &TypedCommand,
        select_cmd: &TypedSelectCommand,
    ) -> Result<Self, Rc<TranslationError>> {
        let mut assignments = ordermap::OrderMap::new();
        let mut assigned = HashSet::new();

        for assignment in &select_cmd.projections.assignments {
            let id = assignment.identifier.clone().valid()?;
            // Use the original AST expression
            assignments.insert(id.clone(), assignment.expression.ast.clone());
            // Track as assigned - LET referencing these fields should trigger barrier
            assigned.insert(id);
        }

        Ok(Self {
            assignments,
            assigned,
            span: command.ast.span,
        })
    }

    /// Initialize from a LET command (LET assignments first, then passthrough)
    ///
    /// LET prepends fields to the schema, so we put LET fields first.
    fn from_let(
        command: &TypedCommand,
        let_cmd: &TypedLetCommand,
    ) -> Result<Self, Rc<TranslationError>> {
        let mut assignments = ordermap::OrderMap::new();
        let mut assigned = HashSet::new();

        // First, add LET's assignments (they go to the front)
        for assignment in &let_cmd.projections.assignments {
            let id = assignment.identifier.clone().valid()?;
            assignments.insert(id.clone(), assignment.expression.ast.clone());
            assigned.insert(id);
        }

        // Then add passthrough projections for all input fields (skipping any that LET set)
        let input_struct = command.input_schema.flatten();
        for (field_name, _field_type) in input_struct.fields.iter() {
            let identifier: Identifier = SimpleIdentifier::from(field_name.clone()).into();
            if !assignments.contains_key(&identifier) {
                let passthrough_expr = synthesize_passthrough_ast(&identifier);
                assignments.insert(identifier, Rc::new(passthrough_expr));
            }
        }

        Ok(Self {
            assignments,
            assigned,
            span: command.ast.span,
        })
    }

    /// Initialize from a DROP command (passthrough minus dropped fields)
    fn from_drop(command: &TypedCommand, drop_cmd: &TypedDropCommand) -> Self {
        let mut assignments = ordermap::OrderMap::new();
        let dropped_set: HashSet<_> = drop_cmd.dropped_fields.iter().cloned().collect();

        // Start with passthrough projections for all input fields except dropped
        let input_struct = command.input_schema.flatten();
        for (field_name, _field_type) in input_struct.fields.iter() {
            let identifier: Identifier = SimpleIdentifier::from(field_name.clone()).into();

            // Skip dropped fields
            if dropped_set.contains(&identifier) {
                continue;
            }

            let passthrough_expr = synthesize_passthrough_ast(&identifier);
            assignments.insert(identifier, Rc::new(passthrough_expr));
        }

        Self {
            assignments,
            assigned: HashSet::new(), // DROP doesn't assign anything
            span: command.ast.span,
        }
    }

    /// Merge a LET command into this pending SELECT
    ///
    /// LET prepends fields to the schema, so we need to insert at the front.
    fn merge_let(
        &mut self,
        command: &TypedCommand,
        let_cmd: &TypedLetCommand,
    ) -> Result<(), Rc<TranslationError>> {
        // Build new assignments with LET fields first, then existing fields
        let mut new_assignments = OrderMap::new();

        // First, add LET's assignments (they go to the front)
        for assignment in &let_cmd.projections.assignments {
            let id = assignment.identifier.clone().valid()?;
            new_assignments.insert(id.clone(), assignment.expression.ast.clone());
            self.assigned.insert(id);
        }

        // Then add existing assignments (skipping any that LET overwrote)
        for (id, expr) in self.assignments.drain(..) {
            if !new_assignments.contains_key(&id) {
                new_assignments.insert(id, expr);
            }
        }

        self.assignments = new_assignments;
        self.expand_span(&command.ast.span);
        Ok(())
    }

    /// Merge a DROP command into this pending SELECT
    fn merge_drop(&mut self, command: &TypedCommand, drop_cmd: &TypedDropCommand) {
        for dropped in &drop_cmd.dropped_fields {
            self.assignments.remove(dropped);
            self.assigned.remove(dropped);
        }
        self.expand_span(&command.ast.span);
    }

    /// Expand the span to include another span (union of the two ranges)
    fn expand_span(&mut self, other: &Span) {
        if other.is_none() {
            return;
        }
        if self.span.is_none() {
            self.span = other.clone();
            return;
        }

        let self_start = self.span.start().unwrap();
        let self_end = self.span.end().unwrap();
        let other_start = other.start().unwrap();
        let other_end = other.end().unwrap();

        let new_start = self_start.min(other_start);
        let new_end = self_end.max(other_end);
        self.span = Span::new(new_start, new_end);
    }

    /// Emit as an AST SELECT command.
    ///
    /// Emits assignments in order. Compound identifiers are preserved as-is;
    /// the IR conversion layer will pack them into struct literals.
    fn emit(self) -> Command {
        let mut builder = select_command();

        for (identifier, expr) in self.assignments {
            builder = builder.named_field(identifier, expr);
        }

        builder.at(self.span).build()
    }
}

/// Synthesize a passthrough AST expression for an identifier.
///
/// For simple identifiers (x), creates a column reference.
/// For compound identifiers (x.y.z), creates a field lookup chain.
fn synthesize_passthrough_ast(identifier: &Identifier) -> Expression {
    match identifier {
        Identifier::Simple(simple) => column_ref(simple.as_str()).build(),
        Identifier::Compound(compound) => {
            // Build chain: column_ref(first).field(second).field(third)...
            let parts = &compound.parts;
            assert!(!parts.is_empty());

            let mut current: Box<dyn ExpressionBuilder> = Box::new(column_ref(parts[0].as_str()));

            for part in &parts[1..] {
                current = Box::new(field(current, part.as_str()));
            }

            current.build()
        }
    }
}

/// Check if two identifiers overlap (one is prefix of the other, or they're equal).
fn identifiers_overlap(a: &Identifier, b: &Identifier) -> bool {
    a == b || a.has_prefix(b) || b.has_prefix(a)
}

/// Extract all column references from a Projections structure.
fn extract_column_references_from_projections(projections: &Projections) -> HashSet<Identifier> {
    let mut refs = HashSet::new();
    for assignment in &projections.assignments {
        extract_column_references_from_expression(&assignment.expression, &mut refs);
    }
    refs
}

/// Extract all column references from a TypedExpression using find() traversal.
fn extract_column_references_from_expression(
    expr: &TypedExpression,
    refs: &mut HashSet<Identifier>,
) {
    // Use find() to traverse all nodes, collecting column references
    // We return false to keep searching (visit all nodes)
    expr.find(&mut |e| {
        if let TypedExpressionKind::ColumnReference(col_ref) = &e.kind {
            if let Ok(simple) = col_ref.column_name.clone().valid() {
                refs.insert(simple.into());
            }
        }
        false // Continue searching to visit all nodes
    });
}

#[cfg(test)]
mod tests {
    use super::*;
    use hamelin_lib::{
        tree::ast::identifier::CompoundIdentifier,
        tree::{
            ast::{pipeline::Pipeline, IntoTyped, TypeCheckExecutor},
            builder::{add, column_ref, drop_command, let_command, pipeline, select_command},
        },
        types::{struct_type::Struct, INT},
    };
    use pretty_assertions::assert_eq;
    use rstest::rstest;
    use std::rc::Rc;

    #[rstest]
    // Case 1: No LET or DROP commands - passes through unchanged
    #[case::no_fusion_needed(
        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: SELECT + LET → fused to single SELECT (both projections merged)
    #[case::select_let_fused(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(let_command().named_field("b", 2).build())
            .build(),
        pipeline()
            .command(select_command().named_field("b", 2).named_field("a", 1).build())
            .build(),
        Struct::default().with_str("b", INT).with_str("a", INT)
    )]
    // Case 3: SELECT + DROP → fused to single SELECT (field removed)
    #[case::select_drop_fused(
        pipeline()
            .command(select_command().named_field("a", 1).named_field("b", 2).build())
            .command(drop_command().field("b").build())
            .build(),
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .build(),
        Struct::default().with_str("a", INT)
    )]
    // Case 4: SELECT + multiple LETs → single SELECT
    #[case::select_multiple_lets_fused(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(let_command().named_field("b", 2).build())
            .command(let_command().named_field("c", 3).build())
            .build(),
        pipeline()
            .command(select_command()
                .named_field("c", 3)
                .named_field("b", 2)
                .named_field("a", 1)
                .build())
            .build(),
        Struct::default().with_str("c", INT).with_str("b", INT).with_str("a", INT)
    )]
    // Case 5: LET references field assigned by SELECT → barrier
    // SELECT a = 1 | LET b = a + 1
    // The LET references 'a' which SELECT assigned, so they cannot fuse.
    #[case::barrier_let_refs_select_field(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(let_command().named_field("b", add(column_ref("a"), 1)).build())
            .build(),
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(select_command()
                .named_field("b", add(column_ref("a"), 1))
                .named_field("a", column_ref("a"))
                .build())
            .build(),
        Struct::default().with_str("b", INT).with_str("a", INT)
    )]
    // Case 6: LET references field assigned by preceding LET → barrier
    // SELECT a = 1 | LET b = 2 | LET c = b + 1
    // First LET doesn't reference 'a', so it fuses with SELECT.
    // Second LET references 'b' which first LET assigned, so barrier.
    #[case::barrier_let_refs_let_field(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(let_command().named_field("b", 2).build())
            .command(let_command().named_field("c", add(column_ref("b"), 1)).build())
            .build(),
        pipeline()
            .command(select_command()
                .named_field("b", 2)
                .named_field("a", 1)
                .build())
            .command(select_command()
                .named_field("c", add(column_ref("b"), 1))
                .named_field("b", column_ref("b"))
                .named_field("a", column_ref("a"))
                .build())
            .build(),
        Struct::default().with_str("c", INT).with_str("b", INT).with_str("a", INT)
    )]
    // Case 7: Chained LET dependencies → multiple barriers
    // SELECT a = 1 | LET b = a + 1 | LET c = b + 1
    // Each LET references the field assigned by the previous command.
    #[case::barrier_chained_dependencies(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(let_command().named_field("b", add(column_ref("a"), 1)).build())
            .command(let_command().named_field("c", add(column_ref("b"), 1)).build())
            .build(),
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(select_command()
                .named_field("b", add(column_ref("a"), 1))
                .named_field("a", column_ref("a"))
                .build())
            .command(select_command()
                .named_field("c", add(column_ref("b"), 1))
                .named_field("b", column_ref("b"))
                .named_field("a", column_ref("a"))
                .build())
            .build(),
        Struct::default().with_str("c", INT).with_str("b", INT).with_str("a", INT)
    )]
    // Case 8: LET overwrites same field without reference → fuses (last write wins)
    // LET a = 1 | LET a = 2
    // Second LET doesn't reference 'a', it just overwrites it. No barrier needed.
    #[case::no_barrier_overwrite_without_ref(
        pipeline()
            .command(let_command().named_field("a", 1).build())
            .command(let_command().named_field("a", 2).build())
            .build(),
        pipeline()
            // Second LET's assignment wins (last write)
            .command(select_command().named_field("a", 2).build())
            .build(),
        Struct::default().with_str("a", INT)
    )]
    // Case 9: LET self-reference → barrier
    // SELECT a = 1 | LET a = a + 1
    // The second LET references 'a' which SELECT assigned.
    #[case::barrier_self_reference(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(let_command().named_field("a", add(column_ref("a"), 1)).build())
            .build(),
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(select_command().named_field("a", add(column_ref("a"), 1)).build())
            .build(),
        Struct::default().with_str("a", INT)
    )]
    // Case 10: Independent LETs fuse (no dependency)
    // SELECT a = 1 | LET b = 2 | LET c = 3
    // None of the LETs reference assigned fields, so all fuse.
    #[case::no_barrier_independent_lets(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(let_command().named_field("b", 2).build())
            .command(let_command().named_field("c", 3).build())
            .build(),
        pipeline()
            .command(select_command()
                .named_field("c", 3)
                .named_field("b", 2)
                .named_field("a", 1)
                .build())
            .build(),
        Struct::default().with_str("c", INT).with_str("b", INT).with_str("a", INT)
    )]
    // Case 11: Three LETs without SELECT - verifies from_let prepend order
    // LET a = 1 | LET b = 2 | LET c = 3
    // First LET uses from_let, subsequent LETs use merge_let.
    // Each successive LET prepends, so final order is c, b, a.
    #[case::three_lets_prepend_order(
        pipeline()
            .command(let_command().named_field("a", 1).build())
            .command(let_command().named_field("b", 2).build())
            .command(let_command().named_field("c", 3).build())
            .build(),
        pipeline()
            .command(select_command()
                .named_field("c", 3)
                .named_field("b", 2)
                .named_field("a", 1)
                .build())
            .build(),
        Struct::default().with_str("c", INT).with_str("b", INT).with_str("a", INT)
    )]
    // Case 12: Compound LET assignments are preserved (struct packing happens in IR)
    // LET x.a = 1 | LET x.b = 2 fuses to SELECT x.b = 2, x.a = 1
    #[case::compound_let_preserved(
        pipeline()
            .command(let_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .build())
            .command(let_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    2,
                )
                .build())
            .build(),
        pipeline()
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    2,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .build())
            .build(),
        Struct::default()
            .with_str("x", Struct::default().with_str("b", INT).with_str("a", INT).into())
    )]
    fn test_fuse_projections(
        #[case] input: Pipeline,
        #[case] expected: Pipeline,
        #[case] expected_output_schema: Struct,
    ) {
        let input_typed = input.typed_with().typed();
        let expected_typed = expected.typed_with().typed();

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

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

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