icydb-core 0.257.2

//! Module: query::plan::expr::aggregate_input
//! Responsibility: aggregate-input canonicalization shared by planner builders and frontend lowering.
//! Does not own: aggregate validation policy, grouped execution wiring, or parser frontends.
//! Boundary: one planner-owned normalization seam for constant folding and numeric literal shaping.

use crate::{
    db::{
        numeric::{apply_value_arithmetic_checked, coerce_numeric_decimal},
        query::plan::{
            AggregateKind,
            expr::{AggregateInputConstantFoldShape, BinaryOp, Expr, Function},
        },
    },
    value::Value,
};

// Keep aggregate input identity canonical anywhere planner-owned aggregate
// expressions are constructed so grouped/global paths do not drift on
// semantically equivalent constant subexpressions.
pub(in crate::db) fn canonicalize_aggregate_input_expr(kind: AggregateKind, expr: Expr) -> Expr {
    let mut folded =
        normalize_aggregate_input_numeric_literals(fold_aggregate_input_constant_expr(expr));

    match kind {
        AggregateKind::Sum | AggregateKind::Avg => {
            if let Expr::Literal(value) = &mut folded
                && let Some(decimal) = coerce_numeric_decimal(value)
            {
                *value = Value::Decimal(decimal.normalize());
            }
            folded
        }
        AggregateKind::Count
        | AggregateKind::Min
        | AggregateKind::Max
        | AggregateKind::Exists
        | AggregateKind::First
        | AggregateKind::Last => folded,
    }
}

// SQL COUNT over any non-null literal is semantically COUNT(*) unless DISTINCT
// is present, so frontends can share the row-count route for those shapes.
pub(in crate::db) const fn aggregate_count_input_expr_is_non_null_literal(expr: &Expr) -> bool {
    matches!(expr, Expr::Literal(value) if !matches!(value, Value::Null))
}

// Fold literal-only aggregate-input subexpressions so aggregate identity
// matching can treat `AVG(age + 1 * 2)` and `AVG(age + 2)` as the same input.
fn fold_aggregate_input_constant_expr(mut expr: Expr) -> Expr {
    expr.map_scalar_children(fold_aggregate_input_constant_expr);
    let folded = match &expr {
        Expr::FunctionCall { function, args } => {
            fold_aggregate_input_constant_function(*function, args)
        }
        Expr::Binary { op, left, right } => fold_aggregate_input_constant_binary(*op, left, right),
        _ => None,
    };

    folded.unwrap_or(expr)
}

// Fold one literal-only binary aggregate-input fragment onto one decimal
// literal so aggregate identity stays stable across equivalent frontend spellings.
fn fold_aggregate_input_constant_binary(op: BinaryOp, left: &Expr, right: &Expr) -> Option<Expr> {
    let (Expr::Literal(left), Expr::Literal(right)) = (left, right) else {
        return None;
    };
    if matches!(left, Value::Null) || matches!(right, Value::Null) {
        return Some(Expr::Literal(Value::Null));
    }

    let arithmetic_op = op.numeric_arithmetic_op()?;
    let result = apply_value_arithmetic_checked(arithmetic_op, left, right)
        .ok()
        .flatten()?;

    Some(Expr::Literal(result))
}

// Fold one admitted literal-only aggregate-input function call when the
// reduced aggregate-input family has one deterministic literal result.
fn fold_aggregate_input_constant_function(function: Function, args: &[Expr]) -> Option<Expr> {
    match function.aggregate_input_constant_fold_shape() {
        Some(AggregateInputConstantFoldShape::Round) => {
            fold_aggregate_input_constant_round(function, args)
        }
        Some(AggregateInputConstantFoldShape::DynamicCoalesce) => {
            fold_aggregate_input_constant_coalesce(args)
        }
        Some(AggregateInputConstantFoldShape::DynamicNullIf) => {
            fold_aggregate_input_constant_nullif(args)
        }
        Some(AggregateInputConstantFoldShape::BinaryNumeric) => {
            fold_aggregate_input_constant_binary_numeric(function, args)
        }
        Some(AggregateInputConstantFoldShape::UnaryNumeric) => {
            fold_aggregate_input_constant_unary_numeric(function, args)
        }
        None => None,
    }
}

// Fold one admitted unary numeric aggregate-input wrapper through the shared
// planner numeric contract so literal-only numeric calls keep one canonical
// aggregate identity.
fn fold_aggregate_input_constant_unary_numeric(function: Function, args: &[Expr]) -> Option<Expr> {
    let [Expr::Literal(input)] = args else {
        return None;
    };
    if matches!(input, Value::Null) {
        return Some(Expr::Literal(Value::Null));
    }

    let decimal = coerce_numeric_decimal(input)?;
    let result = function
        .unary_numeric_function_kind()?
        .eval_decimal(decimal)
        .ok()?;

    Some(Expr::Literal(result))
}

// Fold one admitted binary numeric aggregate-input wrapper through the shared
// planner numeric contract so literal-only numeric calls keep one canonical
// aggregate identity.
fn fold_aggregate_input_constant_binary_numeric(function: Function, args: &[Expr]) -> Option<Expr> {
    let [Expr::Literal(left), Expr::Literal(right)] = args else {
        return None;
    };
    if matches!(left, Value::Null) || matches!(right, Value::Null) {
        return Some(Expr::Literal(Value::Null));
    }

    let result = function
        .binary_numeric_function_kind()?
        .eval_values(left, right)
        .ok()
        .flatten()?;

    Some(Expr::Literal(result))
}

// Fold one literal-only ROUND(...) aggregate-input fragment so parenthesized
// constant arithmetic keeps the same aggregate identity as its literal result.
fn fold_aggregate_input_constant_round(function: Function, args: &[Expr]) -> Option<Expr> {
    let [Expr::Literal(input), Expr::Literal(scale)] = args else {
        return None;
    };
    if matches!(input, Value::Null) || matches!(scale, Value::Null) {
        return Some(Expr::Literal(Value::Null));
    }

    let scale = match scale {
        Value::Int64(value) => u32::try_from(*value).ok()?,
        Value::Nat64(value) => u32::try_from(*value).ok()?,
        _ => return None,
    };

    Some(Expr::Literal(function.eval_numeric_scale(input, scale)?))
}

// Fold one literal-only COALESCE aggregate-input subtree so all-null versus
// first-non-null behavior stays stable before aggregate dedupe compares inputs.
fn fold_aggregate_input_constant_coalesce(args: &[Expr]) -> Option<Expr> {
    let mut literal_values = Vec::with_capacity(args.len());
    for arg in args {
        let Expr::Literal(value) = arg else {
            return None;
        };
        literal_values.push(value.clone());
    }

    Some(Expr::Literal(
        Function::Coalesce.eval_coalesce_values(literal_values.as_slice()),
    ))
}

// Fold one literal-only NULLIF aggregate-input subtree so equivalent frontend
// spellings collapse to the same planner literal before aggregate matching.
fn fold_aggregate_input_constant_nullif(args: &[Expr]) -> Option<Expr> {
    let [Expr::Literal(left), Expr::Literal(right)] = args else {
        return None;
    };

    Some(Expr::Literal(Function::NullIf.eval_nullif_values(
        left,
        right,
        left == right,
    )))
}

// Normalize numeric literal leaves recursively so semantically equivalent
// aggregate inputs like `age + 2` and `age + 1 * 2` share one canonical
// planner identity after literal-only subtree folding.
fn normalize_aggregate_input_numeric_literals(mut expr: Expr) -> Expr {
    expr.map_scalar_children(normalize_aggregate_input_numeric_literals);
    if let Expr::Literal(value) = &mut expr
        && let Some(decimal) = coerce_numeric_decimal(value)
    {
        *value = Value::Decimal(decimal.normalize());
    }

    expr
}