hamelin_translation 0.9.6

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

use std::collections::HashSet;
use std::sync::Arc;

use ordermap::OrderMap;

use hamelin_lib::{
    err::TranslationError,
    tree::{
        ast::{
            command::Command,
            expression::{Expression, ExpressionKind},
            identifier::{CompoundIdentifier, Identifier, SimpleIdentifier},
            node::Span,
        },
        builder::{self, field, field_ref, select_command, ExpressionBuilder},
        typed_ast::{
            clause::Projections,
            command::{
                TypedCommand, TypedCommandKind, TypedDropCommand, TypedSelectCommand,
                TypedSetCommand,
            },
            context::StatementTranslationContext,
            environment::TypeEnvironment,
            expression::{TypedExpression, TypedExpressionKind},
            pipeline::TypedPipeline,
        },
    },
    types::{struct_type::Struct, Type},
};

/// Fuse projection commands (SELECT, SET, DROP) into minimal SELECT commands.
///
/// Pipeline-level pass contract: `Arc<TypedPipeline> -> Result<Arc<TypedPipeline>, ...>`
///
/// This creates a minimal number of SELECT commands by:
/// 1. Accumulating SET/DROP operations into a pending SELECT
/// 2. Detecting dependency barriers when SET 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: Arc<TypedPipeline>,
    ctx: &mut StatementTranslationContext,
) -> Result<Arc<TypedPipeline>, Arc<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(Arc::new(TypedPipeline::from_ast_with_context(
        Arc::new(fused_ast),
        ctx,
    )))
}

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

/// Fuse commands into minimal SELECT commands.
fn fuse_commands(commands: &[Arc<TypedCommand>]) -> Result<Vec<Command>, Arc<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::Set(set_cmd) => {
                let refs = extract_field_references_from_projections(&set_cmd.projections);

                if let Some(ref mut p) = pending {
                    // Check for dependency barrier:
                    // Expression (RHS) references an assigned identifier
                    let expr_depends = refs
                        .iter()
                        .any(|r| p.assigned.iter().any(|a| identifiers_overlap(r, a)));
                    // Target (LHS) is a child of a non-struct-literal assignment
                    // (e.g., SET x.y = 1 when x was assigned as a computed expression).
                    // If the parent assignment is a struct literal, we can modify it in-place.
                    let target_depends = set_cmd.projections.assignments.iter().any(|assignment| {
                        assignment
                            .identifier
                            .valid_ref()
                            .is_ok_and(|target| has_non_literal_ancestor(&p.assignments, target))
                    });
                    if expr_depends || target_depends {
                        if let Some(prev) = pending.take() {
                            result.push(prev.emit()?);
                        }
                        pending = Some(PendingSelect::from_set(command, set_cmd)?);
                    } else {
                        p.merge_set(command, set_cmd)?;
                    }
                } else {
                    pending = Some(PendingSelect::from_set(command, set_cmd)?);
                }
            }

            TypedCommandKind::Drop(drop_cmd) => {
                if let Some(ref mut p) = pending {
                    // DROP is a barrier if it drops a child of a non-struct-literal assignment
                    let drop_depends = drop_cmd
                        .dropped_fields
                        .iter()
                        .any(|dropped| has_non_literal_ancestor(&p.assignments, dropped));
                    if drop_depends {
                        if let Some(prev) = pending.take() {
                            result.push(prev.emit()?);
                        }
                        pending = Some(PendingSelect::from_drop(command, drop_cmd));
                    } else {
                        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)
    /// Only contains explicit assignments from SET/SELECT, NOT passthroughs.
    /// Passthroughs are computed at emit() time from output_schema.
    assignments: OrderMap<Identifier, Arc<Expression>>,
    /// Set of identifiers that have been explicitly assigned (not passthrough)
    /// Used for barrier detection (SET referencing assigned field).
    assigned: HashSet<Identifier>,
    /// Set of identifiers that have been dropped.
    /// Used to determine if struct passthroughs need to recurse.
    dropped: HashSet<Identifier>,
    /// Output schema after all fused commands have been applied.
    /// Used to generate passthroughs at emit() time.
    output_schema: Arc<TypeEnvironment>,
    /// AST span for the synthesized command
    span: Span,
}

impl PendingSelect {
    /// Initialize from a SELECT command
    ///
    /// SELECT replaces the schema entirely, so we store all its assignments
    /// and use its output_schema directly.
    fn from_select(
        command: &TypedCommand,
        select_cmd: &TypedSelectCommand,
    ) -> Result<Self, Arc<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 - SET referencing these fields should trigger barrier
            assigned.insert(id);
        }

        Ok(Self {
            assignments,
            assigned,
            dropped: HashSet::new(),
            output_schema: command.output_schema.clone(),
            span: command.ast.span,
        })
    }

    /// Initialize from a SET command (SET assignments first, then passthrough)
    ///
    /// SET assignment keys appear before existing fields in the output schema,
    /// so we store them first in the OrderMap.
    /// Passthroughs are NOT computed here - they're generated at emit() time
    /// based on output_schema minus assignments.
    fn from_set(
        command: &TypedCommand,
        set_cmd: &TypedSetCommand,
    ) -> Result<Self, Arc<TranslationError>> {
        let mut assignments = ordermap::OrderMap::new();
        let mut assigned = HashSet::new();

        // Store only the SET's explicit assignments
        for assignment in &set_cmd.projections.assignments {
            let id = assignment.identifier.clone().valid()?;
            assignments.insert(id.clone(), assignment.expression.ast.clone());
            assigned.insert(id);
        }

        Ok(Self {
            assignments,
            assigned,
            dropped: HashSet::new(),
            output_schema: command.output_schema.clone(),
            span: command.ast.span,
        })
    }

    /// Initialize from a DROP command (passthrough minus dropped fields)
    ///
    /// DROP doesn't have any explicit assignments - all fields are passthroughs
    /// except for the dropped ones. Passthroughs are generated at emit() time.
    fn from_drop(command: &TypedCommand, drop_cmd: &TypedDropCommand) -> Self {
        // DROP has no explicit assignments - everything comes from output_schema
        // Track dropped fields so struct passthroughs know to recurse
        let dropped = drop_cmd.dropped_fields.iter().cloned().collect();
        Self {
            assignments: ordermap::OrderMap::new(),
            assigned: HashSet::new(),
            dropped,
            output_schema: command.output_schema.clone(),
            span: command.ast.span,
        }
    }

    /// Merge a SET command into this pending SELECT.
    ///
    /// SET assignment keys appear before existing fields in the output schema,
    /// so new assignments are inserted before existing ones in the OrderMap.
    /// Passthroughs are computed at emit() time from the output schema.
    ///
    /// When a compound target's parent is a struct literal in assignments,
    /// we fold the assignment into the literal in-place rather than adding
    /// a separate entry.
    fn merge_set(
        &mut self,
        command: &TypedCommand,
        set_cmd: &TypedSetCommand,
    ) -> Result<(), Arc<TranslationError>> {
        // Resolve all identifiers upfront
        let resolved: Vec<_> = set_cmd
            .projections
            .assignments
            .iter()
            .map(|a| a.identifier.clone().valid().map(|id| (id, a)))
            .collect::<Result<_, _>>()?;

        // First, handle compound targets that modify existing struct literal assignments.
        // These get folded into the literal rather than added as separate entries.
        let mut folded_into_literal = HashSet::new();
        for (id, assignment) in &resolved {
            if let Some((ancestor_id, child_path)) =
                find_struct_literal_ancestor(&self.assignments, id)
            {
                if let Some(ancestor_expr) = self.assignments.get_mut(&ancestor_id) {
                    let folded = modify_struct_literal_field(
                        ancestor_expr,
                        &child_path,
                        Some(assignment.expression.ast.clone()),
                    );
                    if folded {
                        folded_into_literal.insert(id.clone());
                    }
                }
            }
        }

        // Build new assignments with SET fields first, then existing fields
        let mut new_assignments = OrderMap::new();

        // Add SET's non-folded assignments (they go to the front)
        for (id, assignment) in &resolved {
            if folded_into_literal.contains(id) {
                continue;
            }
            new_assignments.insert(id.clone(), assignment.expression.ast.clone());
            self.assigned.insert(id.clone());
        }

        // Then add existing explicit assignments (skipping any that the new SET shadows)
        // If SET assigns x.y and we had x = expr, the new SET shadows our x.
        // If SET assigns x and we had x.y = expr, the new SET also shadows our x.y.
        for (id, expr) in self.assignments.drain(..) {
            if new_assignments
                .keys()
                .any(|new_id| identifiers_overlap(&id, new_id))
            {
                continue;
            }
            new_assignments.insert(id, expr);
        }

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

    /// Merge a DROP command into this pending SELECT
    ///
    /// DROP removes fields. We remove any explicit assignments for dropped fields,
    /// and update output_schema. Passthroughs will be computed at emit() time.
    ///
    /// When dropping a child of a struct literal assignment, we prune the field
    /// from the literal rather than tracking it as a separate drop.
    fn merge_drop(&mut self, command: &TypedCommand, drop_cmd: &TypedDropCommand) {
        for dropped in &drop_cmd.dropped_fields {
            // Try to prune from a struct literal ancestor first
            let pruned = find_struct_literal_ancestor(&self.assignments, dropped)
                .and_then(|(ancestor_id, child_path)| {
                    let expr = self.assignments.get_mut(&ancestor_id)?;
                    modify_struct_literal_field(expr, &child_path, None).then_some(())
                })
                .is_some();

            if !pruned {
                self.assignments.remove(dropped);
                self.assigned.remove(dropped);
                self.dropped.insert(dropped.clone());
            }
        }
        self.output_schema = command.output_schema.clone();
        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;
        }

        if let (Some(self_start), Some(self_end), Some(other_start), Some(other_end)) = (
            self.span.start(),
            self.span.end(),
            other.start(),
            other.end(),
        ) {
            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 fields in the order they appear in `output_schema.to_struct()`:
    /// - Top-level: newest fields first (left-most position)
    /// - Nested structs: oldest fields first (preserves original struct order)
    ///
    /// For each field, either emit the explicit assignment or generate a passthrough.
    fn emit(self) -> Result<Command, Arc<TranslationError>> {
        let mut builder = select_command();

        let output_struct = self.output_schema.as_struct();
        let assignment_set: HashSet<_> = self.assignments.keys().cloned().collect();

        let mut fields_to_emit = Vec::new();
        emit_fields_in_schema_order(
            &output_struct,
            &self.assignments,
            &assignment_set,
            &self.dropped,
            &[],
            &mut fields_to_emit,
        )?;

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

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

/// Check if any assignment or drop overlaps with a struct field.
///
/// Returns true if we need to recurse into the struct to handle partial coverage.
/// Returns false if we can emit the struct directly as a passthrough.
fn any_modification_overlaps_struct(
    assignments: &HashSet<Identifier>,
    dropped: &HashSet<Identifier>,
    struct_id: &Identifier,
) -> bool {
    let overlaps_assignments = assignments
        .iter()
        .any(|a| a == struct_id || a.has_prefix(struct_id) || struct_id.has_prefix(a));
    let overlaps_dropped = dropped
        .iter()
        .any(|d| d == struct_id || d.has_prefix(struct_id) || struct_id.has_prefix(d));
    overlaps_assignments || overlaps_dropped
}

/// Recursively emit fields in schema order, using assignments or passthroughs.
///
/// This preserves the correct field ordering from output_schema.to_struct():
/// - Top-level: newest fields first (left-most position)
/// - Nested structs: oldest fields first (preserves original struct order)
///
/// For each field position:
/// - If there's an explicit assignment → emit that
/// - If it's a struct with partial modifications → recurse for ordered leaf fields
/// - If it's a struct with no modifications → emit struct passthrough
/// - Otherwise → emit leaf passthrough
fn emit_fields_in_schema_order(
    schema: &Struct,
    assignments: &OrderMap<Identifier, Arc<Expression>>,
    assignment_set: &HashSet<Identifier>,
    dropped: &HashSet<Identifier>,
    prefix: &[SimpleIdentifier],
    output: &mut Vec<(Identifier, Arc<Expression>)>,
) -> Result<(), Arc<TranslationError>> {
    for (field_name, field_type) in schema.iter() {
        // Build the full identifier path
        let mut path = prefix.to_vec();
        path.push(field_name.clone().into());
        let full_id: Identifier = match path.as_slice() {
            [] => {
                return Err(TranslationError::fatal(
                    "fuse_projections",
                    "emit_fields_in_schema_order called with empty path".into(),
                )
                .into())
            }
            [single] => single.clone().into(),
            [first, second, rest @ ..] => {
                CompoundIdentifier::new(first.clone(), second.clone(), rest.to_vec()).into()
            }
        };

        // Check if this field is exactly covered by an assignment
        if let Some(expr) = assignments.get(&full_id) {
            output.push((full_id, expr.clone()));
            continue;
        }

        // For struct types, check if any modification overlaps with this struct
        if let Type::Struct(inner_struct) = field_type {
            if any_modification_overlaps_struct(assignment_set, dropped, &full_id) {
                // Partial coverage - recurse to emit leaf-level fields in order
                emit_fields_in_schema_order(
                    inner_struct,
                    assignments,
                    assignment_set,
                    dropped,
                    &path,
                    output,
                )?;
            } else {
                // No overlap - emit the struct directly as a passthrough
                let passthrough_expr = synthesize_passthrough_ast(&full_id);
                output.push((full_id, passthrough_expr.into()));
            }
            continue;
        }

        // Leaf field not covered - emit passthrough
        let passthrough_expr = synthesize_passthrough_ast(&full_id);
        output.push((full_id, passthrough_expr.into()));
    }
    Ok(())
}

/// 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) => field_ref(simple.as_str()).build(),
        Identifier::Compound(compound) => {
            // Build chain: field_ref(first).field(second).field(third)...
            let mut current: Box<dyn ExpressionBuilder> =
                Box::new(field_ref(compound.first().as_str()));

            for part in compound.rest_parts() {
                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)
}

/// Check if an identifier has an ancestor in assignments that can't be modified in-place.
///
/// Returns true if there's a strict prefix of `target` in `assignments` whose expression
/// is not a struct literal, OR is a struct literal but the path through it hits a
/// non-literal intermediate field. This means we can't fold the assignment and need a barrier.
fn has_non_literal_ancestor(
    assignments: &OrderMap<Identifier, Arc<Expression>>,
    target: &Identifier,
) -> bool {
    let segments = target.segments();
    for prefix in target.prefixes() {
        if let Some(expr) = assignments.get(&prefix) {
            let ExpressionKind::StructLiteral(ref lit) = expr.kind else {
                return true;
            };
            // Walk the remaining intermediate segments (excluding the leaf)
            // through nested struct literals. If any intermediate is not a
            // struct literal, we can't fold — barrier needed.
            let prefix_len = prefix.segments().len();
            let remaining = &segments[prefix_len..];
            let mut current = lit;
            for segment in &remaining[..remaining.len().saturating_sub(1)] {
                match current
                    .fields
                    .iter()
                    .find(|(id, _)| id.valid_ref().is_ok_and(|n| n == segment))
                {
                    Some((_, child_expr)) => match &child_expr.kind {
                        ExpressionKind::StructLiteral(nested) => current = nested,
                        _ => return true,
                    },
                    None => return false, // field doesn't exist yet, will be created
                }
            }
            return false;
        }
    }
    false
}

/// Find the nearest ancestor in assignments that is a struct literal.
///
/// Returns `Some((ancestor_id, child_path))` where `child_path` is the segments
/// below the ancestor. E.g., for target `x.a.b` with ancestor `x`, returns
/// `(x, [a, b])`.
fn find_struct_literal_ancestor(
    assignments: &OrderMap<Identifier, Arc<Expression>>,
    target: &Identifier,
) -> Option<(Identifier, Vec<SimpleIdentifier>)> {
    let segments = target.segments();
    for prefix in target.prefixes() {
        if let Some(expr) = assignments.get(&prefix) {
            if matches!(expr.kind, ExpressionKind::StructLiteral(_)) {
                let prefix_len = prefix.segments().len();
                let child_path = segments[prefix_len..].to_vec();
                return Some((prefix, child_path));
            }
        }
    }
    None
}

/// Modify a field within a struct literal expression.
///
/// `child_path` is the path of segments below the struct literal to the target field.
/// `new_value` is `Some(expr)` to set/replace the field, or `None` to remove it.
///
/// For nested paths (e.g., `[a, b]`), navigates into nested struct literals.
/// Returns `true` if the modification succeeded, `false` if the path couldn't be followed.
fn modify_struct_literal_field(
    expr: &mut Arc<Expression>,
    child_path: &[SimpleIdentifier],
    new_value: Option<Arc<Expression>>,
) -> bool {
    if child_path.is_empty() {
        return false;
    }

    // Check feasibility before cloning
    if !matches!(expr.kind, ExpressionKind::StructLiteral(_)) {
        return false;
    }

    // Navigate to the parent struct literal containing the target field.
    // For path [a, b, c]: navigate through a, then b, to reach c's parent.
    let owned = Arc::make_mut(expr);
    let ExpressionKind::StructLiteral(ref mut lit) = owned.kind else {
        return false;
    };
    let mut current = lit;

    for segment in &child_path[..child_path.len() - 1] {
        let Some((_, child_expr)) = current.find_field_mut(segment.as_str()) else {
            return false;
        };
        let child = Arc::make_mut(child_expr);
        let ExpressionKind::StructLiteral(ref mut nested) = child.kind else {
            return false;
        };
        current = nested;
    }

    // Apply the modification to the leaf field
    let leaf = child_path.last().expect("child_path is non-empty").as_str();
    match new_value {
        Some(value) => current.set_field(leaf, value),
        None => current.remove_field(leaf),
    }
    true
}

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

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

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

    #[rstest]
    // Case 1: No SET 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 + SET → fused to single SELECT (both projections merged)
    #[case::select_set_fused(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(set_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 SETs → single SELECT
    #[case::select_multiple_sets_fused(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(set_command().named_field("b", 2).build())
            .command(set_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: SET references field assigned by SELECT → barrier
    // SELECT a = 1 | SET b = a + 1
    // The SET references 'a' which SELECT assigned, so they cannot fuse.
    #[case::barrier_set_refs_select_field(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(set_command().named_field("b", add(field_ref("a"), 1)).build())
            .build(),
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(select_command()
                .named_field("b", add(field_ref("a"), 1))
                .named_field("a", field_ref("a"))
                .build())
            .build(),
        Struct::default().with_str("b", INT).with_str("a", INT)
    )]
    // Case 6: SET references field assigned by preceding SET → barrier
    // SELECT a = 1 | SET b = 2 | SET c = b + 1
    // First SET doesn't reference 'a', so it fuses with SELECT.
    // Second SET references 'b' which first SET assigned, so barrier.
    #[case::barrier_set_refs_set_field(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(set_command().named_field("b", 2).build())
            .command(set_command().named_field("c", add(field_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(field_ref("b"), 1))
                .named_field("b", field_ref("b"))
                .named_field("a", field_ref("a"))
                .build())
            .build(),
        Struct::default().with_str("c", INT).with_str("b", INT).with_str("a", INT)
    )]
    // Case 7: Chained SET dependencies → multiple barriers
    // SELECT a = 1 | SET b = a + 1 | SET c = b + 1
    // Each SET references the field assigned by the previous command.
    #[case::barrier_chained_dependencies(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(set_command().named_field("b", add(field_ref("a"), 1)).build())
            .command(set_command().named_field("c", add(field_ref("b"), 1)).build())
            .build(),
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(select_command()
                .named_field("b", add(field_ref("a"), 1))
                .named_field("a", field_ref("a"))
                .build())
            .command(select_command()
                .named_field("c", add(field_ref("b"), 1))
                .named_field("b", field_ref("b"))
                .named_field("a", field_ref("a"))
                .build())
            .build(),
        Struct::default().with_str("c", INT).with_str("b", INT).with_str("a", INT)
    )]
    // Case 8: SET overwrites same field without reference → fuses (last write wins)
    // SET a = 1 | SET a = 2
    // Second SET doesn't reference 'a', it just overwrites it. No barrier needed.
    #[case::no_barrier_overwrite_without_ref(
        pipeline()
            .command(set_command().named_field("a", 1).build())
            .command(set_command().named_field("a", 2).build())
            .build(),
        pipeline()
            // Second SET's assignment wins (last write)
            .command(select_command().named_field("a", 2).build())
            .build(),
        Struct::default().with_str("a", INT)
    )]
    // Case 9: SET self-reference → barrier
    // SELECT a = 1 | SET a = a + 1
    // The second SET references 'a' which SELECT assigned.
    #[case::barrier_self_reference(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(set_command().named_field("a", add(field_ref("a"), 1)).build())
            .build(),
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(select_command().named_field("a", add(field_ref("a"), 1)).build())
            .build(),
        Struct::default().with_str("a", INT)
    )]
    // Case 10: Independent SETs fuse (no dependency)
    // SELECT a = 1 | SET b = 2 | SET c = 3
    // None of the SETs reference assigned fields, so all fuse.
    #[case::no_barrier_independent_sets(
        pipeline()
            .command(select_command().named_field("a", 1).build())
            .command(set_command().named_field("b", 2).build())
            .command(set_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 SETs without SELECT - verifies from_set prepend order
    // SET a = 1 | SET b = 2 | SET c = 3
    // First SET uses from_set, subsequent SETs use merge_set.
    // Each successive SET prepends, so final order is c, b, a.
    #[case::three_sets_prepend_order(
        pipeline()
            .command(set_command().named_field("a", 1).build())
            .command(set_command().named_field("b", 2).build())
            .command(set_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 SET assignments are preserved (struct packing happens in IR)
    // SET x.a = 1 | SET x.b = 2 fuses to SELECT x.a = 1, x.b = 2
    // Schema order: a first (created by first SET), b appended by second SET.
    #[case::compound_set_preserved(
        pipeline()
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .build())
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    2,
                )
                .build())
            .build(),
        pipeline()
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    2,
                )
                .build())
            .build(),
        Struct::default()
            .with_str("x", Struct::default().with_str("a", INT).with_str("b", INT).into())
    )]
    // Case 13: Compound SET after EXPLODE barrier with nested struct in input schema
    // This is the exact scenario from the EXPLODE normalization bug:
    // After EXPLODE, input schema has {temp: Int, data: Struct{arr: Array<Int>}}
    // Then SET data.arr = temp assigns a compound identifier.
    // Bug: from_set was adding passthrough for `data` (simple) from flatten() output,
    // which conflicts with the compound assignment `data.arr`.
    // Fix: from_set must skip passthroughs that overlap with compound assignments.
    #[case::compound_set_after_explode_barrier(
        pipeline()
            // Setup: create nested struct and temp field
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("data".into(), "arr".into(), vec![]),
                    array().element(1).element(2).element(3),
                )
                .named_field("temp", 42)
                .build())
            // EXPLODE is a barrier
            .command(hamelin_lib::tree::builder::explode_command()
                .named_field("temp", field_ref("temp"))
                .build())
            // After barrier: SET with compound identifier into existing struct
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("data".into(), "arr".into(), vec![]),
                    field_ref("temp"),
                )
                .build())
            .build(),
        pipeline()
            // First SELECT (before barrier)
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("data".into(), "arr".into(), vec![]),
                    array().element(1).element(2).element(3),
                )
                .named_field("temp", 42)
                .build())
            // EXPLODE barrier preserved
            .command(hamelin_lib::tree::builder::explode_command()
                .named_field("temp", field_ref("temp"))
                .build())
            // Second SELECT: compound assignment + temp passthrough, but NO `data` passthrough
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("data".into(), "arr".into(), vec![]),
                    field_ref("temp"),
                )
                .named_field("temp", field_ref("temp"))
                .build())
            .build(),
        Struct::default()
            .with_str("data", Struct::default().with_str("arr", INT).into())
            .with_str("temp", INT)
    )]
    // Case 14: Compound SET adding sibling field - must preserve existing siblings
    // SET x.a = 1, x.b = 2 | WHERE true | SET x.y = 3
    // After barrier, input schema has {x: Struct{a: Int, b: Int}}.
    // The SET assigns x.y, which appends to end of x's struct.
    // Expected output schema: {x: Struct{a: Int, b: Int, y: Int}}
    #[case::compound_set_sibling_preserves_existing(
        pipeline()
            // Setup: create struct with two fields
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    2,
                )
                .build())
            // WHERE is a barrier
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // After barrier: SET adds a new sibling field to existing struct
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "y".into(), vec![]),
                    3,
                )
                .build())
            .build(),
        pipeline()
            // First SELECT (before barrier)
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    2,
                )
                .build())
            // WHERE barrier preserved
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // Second SELECT: passthroughs in schema order, then new field at end
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    hamelin_lib::tree::builder::field(field_ref("x"), "a"),
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    hamelin_lib::tree::builder::field(field_ref("x"), "b"),
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "y".into(), vec![]),
                    3,
                )
                .build())
            .build(),
        Struct::default()
            .with_str("x", Struct::default().with_str("a", INT).with_str("b", INT).with_str("y", INT).into())
    )]
    // Case 15: Deep sibling preservation
    // SET x.c.d = 4, x.c.e = 5 | WHERE true | SET x.c.d = 6
    // Expect x.c.e preserved, x.c.d overwritten.
    #[case::deep_sibling_preservation(
        pipeline()
            // Setup: create deeply nested struct with two fields
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["d".into()]),
                    4,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["e".into()]),
                    5,
                )
                .build())
            // WHERE is a barrier
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // After barrier: SET overwrites one deep field
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["d".into()]),
                    6,
                )
                .build())
            .build(),
        pipeline()
            // First SELECT (before barrier)
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["d".into()]),
                    4,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["e".into()]),
                    5,
                )
                .build())
            // WHERE barrier preserved
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // Second SELECT: overwritten field + passthrough for sibling
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["d".into()]),
                    6,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["e".into()]),
                    hamelin_lib::tree::builder::field(
                        hamelin_lib::tree::builder::field(field_ref("x"), "c"),
                        "e"
                    ),
                )
                .build())
            .build(),
        Struct::default()
            .with_str("x", Struct::default()
                .with_str("c", Struct::default()
                    .with_str("d", INT)
                    .with_str("e", INT)
                    .into())
                .into())
    )]
    // Case 16: Parent vs child conflict - parent first, then child assignment
    // SET x = {a: 1, b: 2} | WHERE true | SET x.a = 3
    // Expect x.a overwritten, x.b preserved via leaf passthrough.
    // First SELECT emits struct literal directly (no desugaring).
    #[case::parent_then_child_assignment(
        pipeline()
            // Setup: create struct via struct literal
            .command(set_command()
                .named_field(
                    "x",
                    hamelin_lib::tree::builder::struct_literal()
                        .field("a", 1)
                        .field("b", 2),
                )
                .build())
            // WHERE is a barrier
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // After barrier: SET assigns a child field
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    3,
                )
                .build())
            .build(),
        pipeline()
            // First SELECT: struct literal emitted directly
            .command(select_command()
                .named_field(
                    "x",
                    hamelin_lib::tree::builder::struct_literal()
                        .field("a", 1)
                        .field("b", 2),
                )
                .build())
            // WHERE barrier preserved
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // Second SELECT: child assignment + leaf passthrough for sibling
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    3,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    hamelin_lib::tree::builder::field(field_ref("x"), "b"),
                )
                .build())
            .build(),
        Struct::default()
            .with_str("x", Struct::default().with_str("a", INT).with_str("b", INT).into())
    )]
    // Case 17: Child vs parent conflict - child first, then parent assignment
    // SET x.a = 1 | WHERE true | SET x = {y: 2}
    // Parent struct literal REPLACES x entirely. Struct literal emitted directly.
    #[case::child_then_parent_assignment(
        pipeline()
            // Setup: create child field
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .build())
            // WHERE is a barrier
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // After barrier: SET assigns struct literal - replaces x entirely
            .command(set_command()
                .named_field(
                    "x",
                    hamelin_lib::tree::builder::struct_literal().field("y", 2),
                )
                .build())
            .build(),
        pipeline()
            // First SELECT (before barrier)
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .build())
            // WHERE barrier preserved
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // Second SELECT: struct literal emitted directly
            .command(select_command()
                .named_field(
                    "x",
                    hamelin_lib::tree::builder::struct_literal().field("y", 2),
                )
                .build())
            .build(),
        Struct::default()
            .with_str("x", Struct::default().with_str("y", INT).into())
    )]
    // Case 18: Mixed deep + new field (recursive example)
    // Input schema x: {a, b, c: {d, e}}, SET assigns x.y = 3, x.c.d = 4
    // Schema order: x: {a, b, c: {d, e}, y} — y appended at end, d replaced in-place.
    #[case::mixed_deep_and_new_field(
        pipeline()
            // Setup: create struct with nested structure
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    2,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["d".into()]),
                    10,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["e".into()]),
                    20,
                )
                .build())
            // WHERE is a barrier
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // After barrier: SET assigns new sibling + overwrites deep field
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "y".into(), vec![]),
                    3,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["d".into()]),
                    4,
                )
                .build())
            .build(),
        pipeline()
            // First SELECT (before barrier)
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    2,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["d".into()]),
                    10,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["e".into()]),
                    20,
                )
                .build())
            // WHERE barrier preserved
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // Second SELECT: schema order x.a, x.b, x.c.d, x.c.e, x.y
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    hamelin_lib::tree::builder::field(field_ref("x"), "a"),
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    hamelin_lib::tree::builder::field(field_ref("x"), "b"),
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["d".into()]),
                    4,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec!["e".into()]),
                    hamelin_lib::tree::builder::field(
                        hamelin_lib::tree::builder::field(field_ref("x"), "c"),
                        "e"
                    ),
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "y".into(), vec![]),
                    3,
                )
                .build())
            .build(),
        Struct::default()
            .with_str("x", Struct::default()
                .with_str("a", INT)
                .with_str("b", INT)
                .with_str("c", Struct::default().with_str("d", INT).with_str("e", INT).into())
                .with_str("y", INT)
                .into())
    )]
    // Case 19: DROP + SET within same fused block
    // SET x.a = 1, x.b = 2, x.c = 3 | WHERE true | DROP x.b | SET x.d = 4
    // Schema order: x: {a, c, d} — b dropped, d appended at end.
    #[case::drop_and_set_fused(
        pipeline()
            // Setup: create struct with three fields
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    2,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec![]),
                    3,
                )
                .build())
            // WHERE is a barrier
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // After barrier: DROP one field, then SET adds another
            .command(drop_command()
                .field(CompoundIdentifier::new("x".into(), "b".into(), vec![]))
                .build())
            .command(set_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "d".into(), vec![]),
                    4,
                )
                .build())
            .build(),
        pipeline()
            // First SELECT (before barrier)
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    1,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    2,
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec![]),
                    3,
                )
                .build())
            // WHERE barrier preserved
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // Second SELECT: schema order x.a, x.c, x.d
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "a".into(), vec![]),
                    hamelin_lib::tree::builder::field(field_ref("x"), "a"),
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec![]),
                    hamelin_lib::tree::builder::field(field_ref("x"), "c"),
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "d".into(), vec![]),
                    4,
                )
                .build())
            .build(),
        Struct::default()
            .with_str("x", Struct::default()
                .with_str("a", INT)
                .with_str("c", INT)
                .with_str("d", INT)
                .into())
    )]
    // Case 20: DROP child from direct struct binding
    // SET x = {a: 1, b: 2, c: 3} | WHERE true | DROP x.a
    // First SELECT emits struct literal directly. After barrier, DROP is a passthrough case.
    #[case::drop_child_from_struct_binding(
        pipeline()
            // Setup: create struct via struct literal
            .command(set_command()
                .named_field(
                    "x",
                    struct_literal()
                        .field("a", 1)
                        .field("b", 2)
                        .field("c", 3),
                )
                .build())
            // WHERE is a barrier
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // After barrier: DROP a child field
            .command(drop_command()
                .field(CompoundIdentifier::new("x".into(), "a".into(), vec![]))
                .build())
            .build(),
        pipeline()
            // First SELECT: struct literal emitted directly
            .command(select_command()
                .named_field(
                    "x",
                    struct_literal()
                        .field("a", 1)
                        .field("b", 2)
                        .field("c", 3),
                )
                .build())
            // WHERE barrier preserved
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // Second SELECT: passthroughs for remaining children (x.a dropped)
            .command(select_command()
                .named_field(
                    CompoundIdentifier::new("x".into(), "b".into(), vec![]),
                    hamelin_lib::tree::builder::field(field_ref("x"), "b"),
                )
                .named_field(
                    CompoundIdentifier::new("x".into(), "c".into(), vec![]),
                    hamelin_lib::tree::builder::field(field_ref("x"), "c"),
                )
                .build())
            .build(),
        Struct::default()
            .with_str("x", Struct::default().with_str("b", INT).with_str("c", INT).into())
    )]
    // Case 21: No overlap sanity - unrelated fields passthrough unchanged
    // SET a = 1 | WHERE true | SET b = 2
    // Expect both fields present, no interference.
    #[case::no_overlap_unrelated_fields(
        pipeline()
            // Setup: create one field
            .command(set_command()
                .named_field("a", 1)
                .build())
            // WHERE is a barrier
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // After barrier: SET adds unrelated field
            .command(set_command()
                .named_field("b", 2)
                .build())
            .build(),
        pipeline()
            // First SELECT (before barrier)
            .command(select_command()
                .named_field("a", 1)
                .build())
            // WHERE barrier preserved
            .command(hamelin_lib::tree::builder::where_command(true).build())
            // Second SELECT: new field + passthrough for existing
            .command(select_command()
                .named_field("b", 2)
                .named_field("a", field_ref("a"))
                .build())
            .build(),
        Struct::default().with_str("b", INT).with_str("a", INT)
    )]
    fn test_fuse_projections(
        #[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 = fuse_projections(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);
    }
}