icydb-core 0.264.3

use crate::db::query::preparation::PreparationWork;
use crate::db::{
    query::plan::{
        GroupField,
        expr::{Alias, Expr, FieldId, ProjectionField, ProjectionSelection},
    },
    schema::SchemaInfo,
    sql::{
        identifier::split_qualified_identifier,
        lowering::{
            AnalyzedLoweredExpr, LoweredExprAnalysis, SqlLoweringError,
            aggregate::SqlAggregateCallInterner,
            expr::{SqlExprPhase, lower_sql_expr},
            select::order::LoweredSqlOrderTerm,
        },
        parser::{SqlAggregateCall, SqlProjection, SqlSelectItem},
    },
};
use icydb_diagnostic_code::{DiagnosticExecutionBudgetResource as Resource, QueryFieldRole};

///
/// LoweredSqlProjectionSelection
///
/// SQL-local projection wrapper that keeps expression analyses beside lowered
/// projection fields until schema-bound SQL capability validation consumes
/// them.
///

#[derive(Clone, Debug)]
pub(super) struct LoweredSqlProjectionSelection {
    selection: ProjectionSelection,
    expr_analyses: Vec<LoweredExprAnalysis>,
}

impl LoweredSqlProjectionSelection {
    #[must_use]
    pub(super) const fn all() -> Self {
        Self {
            selection: ProjectionSelection::All,
            expr_analyses: Vec::new(),
        }
    }

    #[must_use]
    pub(super) const fn fields(fields: Vec<FieldId>) -> Self {
        Self {
            selection: ProjectionSelection::Fields(fields),
            expr_analyses: Vec::new(),
        }
    }

    #[must_use]
    pub(super) fn exprs(fields: Vec<ProjectionField>, analyses: Vec<LoweredExprAnalysis>) -> Self {
        debug_assert_eq!(
            fields.len(),
            analyses.len(),
            "lowered SQL projection analysis must stay aligned with projection fields",
        );
        Self {
            selection: ProjectionSelection::Exprs(fields),
            expr_analyses: analyses,
        }
    }

    #[must_use]
    pub(super) const fn selection(&self) -> &ProjectionSelection {
        &self.selection
    }

    #[must_use]
    pub(super) const fn expr_analyses(&self) -> &[LoweredExprAnalysis] {
        self.expr_analyses.as_slice()
    }

    #[must_use]
    pub(super) fn into_selection(self) -> ProjectionSelection {
        self.selection
    }

    fn is_grouped_canonical_identity(&self, group_by: &[String]) -> bool {
        matches!(
            &self.selection,
            ProjectionSelection::Exprs(fields)
                if grouped_projection_is_canonical_identity(fields.as_slice(), group_by)
        )
    }
}

///
/// LoweredGroupedProjection
///
/// SQL-local grouped projection artifact that keeps the projection and the
/// first-seen unique aggregate calls from the same analyzed expression pass.
///

#[derive(Clone, Debug)]
pub(super) struct LoweredGroupedProjection {
    selection: LoweredSqlProjectionSelection,
    aggregate_calls: Vec<SqlAggregateCall>,
}

impl LoweredGroupedProjection {
    #[must_use]
    pub(super) const fn aggregate_calls(&self) -> &[SqlAggregateCall] {
        self.aggregate_calls.as_slice()
    }

    #[must_use]
    pub(super) fn into_projection_selection(
        self,
        allow_identity_fast_path: bool,
        group_by: &[String],
    ) -> LoweredSqlProjectionSelection {
        if allow_identity_fast_path && self.selection.is_grouped_canonical_identity(group_by) {
            LoweredSqlProjectionSelection::all()
        } else {
            self.selection
        }
    }
}

pub(super) fn lower_scalar_projection_selection(
    projection: SqlProjection,
    projection_aliases: &[Option<String>],
    work: &PreparationWork<'_>,
) -> Result<LoweredSqlProjectionSelection, SqlLoweringError> {
    let SqlProjection::Items(items) = projection else {
        return Ok(LoweredSqlProjectionSelection::all());
    };

    if let Some(field_ids) = direct_scalar_field_selection(items.as_slice(), projection_aliases) {
        return Ok(LoweredSqlProjectionSelection::fields(field_ids));
    }

    let mut fields = Vec::with_capacity(items.len());
    let mut projection_facts = Vec::with_capacity(items.len());
    for (index, item) in items.into_iter().enumerate() {
        let analyzed = lower_analyzed_select_item_expr(&item, SqlExprPhase::Scalar, work)?;
        let (expr, expr_facts) = analyzed.into_parts();
        fields.push(ProjectionField::Scalar {
            expr,
            alias: projection_aliases
                .get(index)
                .and_then(Option::as_deref)
                .map(Alias::new),
        });
        projection_facts.push(expr_facts);
    }

    Ok(LoweredSqlProjectionSelection::exprs(
        fields,
        projection_facts,
    ))
}

pub(super) fn lower_grouped_projection(
    projection: SqlProjection,
    projection_aliases: &[Option<String>],
    group_by: &[String],
    schema: &SchemaInfo,
    work: &PreparationWork<'_>,
) -> Result<LoweredGroupedProjection, SqlLoweringError> {
    if group_by.is_empty() {
        return Err(SqlLoweringError::unsupported_select_group_by());
    }

    let SqlProjection::Items(items) = projection else {
        return Err(SqlLoweringError::grouped_projection_requires_explicit_list());
    };
    // Validate every declared key, even if no projected expression references
    // it. Retained execution keys belong to structural planning, not this proof.
    for field in group_by {
        work.charge(Resource::PredicateExpressionSteps, 1 + field.len() as u64)?;
        if GroupField::accepted_kind_for_label(schema, field).is_none() {
            return Err(SqlLoweringError::unknown_field(
                QueryFieldRole::GroupBy,
                field,
            ));
        }
    }

    let mut seen_aggregate = false;
    let mut fields = Vec::with_capacity(items.len());
    let mut projection_facts = Vec::with_capacity(items.len());
    let mut aggregate_calls = Vec::new();
    let mut aggregate_call_interner = SqlAggregateCallInterner::new();

    for (index, item) in items.into_iter().enumerate() {
        let analyzed = lower_analyzed_select_item_expr(&item, SqlExprPhase::PostAggregate, work)?;
        let expr_facts = analyzed.analysis();
        let contains_aggregate = expr_facts.contains_aggregate();
        if seen_aggregate && !contains_aggregate {
            return Err(SqlLoweringError::grouped_projection_scalar_after_aggregate(
                index,
            ));
        }
        validate_grouped_projection_expr(index, group_by, schema, expr_facts, work)?;
        seen_aggregate |= contains_aggregate;
        if contains_aggregate {
            aggregate_call_interner.extend_select_item(&mut aggregate_calls, &item)?;
        }

        let (expr, expr_facts) = analyzed.into_parts();
        fields.push(ProjectionField::Scalar {
            expr,
            alias: projection_aliases
                .get(index)
                .and_then(Option::as_deref)
                .map(Alias::new),
        });
        projection_facts.push(expr_facts);
    }

    if !seen_aggregate {
        return Err(SqlLoweringError::grouped_projection_requires_aggregate());
    }

    Ok(LoweredGroupedProjection {
        selection: LoweredSqlProjectionSelection::exprs(fields, projection_facts),
        aggregate_calls,
    })
}

// Validate one grouped projection expression against grouped-key authority
// while preserving specific unknown-field diagnostics.
fn validate_grouped_projection_expr(
    index: usize,
    group_fields: &[String],
    schema: &SchemaInfo,
    analysis: &LoweredExprAnalysis,
    work: &PreparationWork<'_>,
) -> Result<(), SqlLoweringError> {
    if let Some(field) = analysis.first_unknown_field_for_schema(schema) {
        return Err(SqlLoweringError::unknown_field(
            QueryFieldRole::Projection,
            field,
        ));
    }
    if !analysis.references_only_group_fields(group_fields, work)? {
        return Err(SqlLoweringError::grouped_projection_references_non_group_field(index));
    }

    Ok(())
}

// Preserve the older grouped `ProjectionSelection::All` fast path only for
// the canonical identity shape where projected grouped fields match `GROUP BY`
// exactly and aggregate terminals follow in declaration order.
fn grouped_projection_is_canonical_identity(
    fields: &[ProjectionField],
    group_by: &[String],
) -> bool {
    if fields.len() < group_by.len() {
        return false;
    }

    let Some((group_fields, aggregate_fields)) = fields.split_at_checked(group_by.len()) else {
        return false;
    };

    group_fields
        .iter()
        .zip(group_by.iter())
        .all(|(field, group_by)| match field {
            ProjectionField::Scalar {
                expr: Expr::Field(field_id),
                alias: None,
            } => field_id.as_str() == group_by,
            ProjectionField::Scalar {
                expr: Expr::FieldPath(path),
                alias: None,
            } => path.path_spec().matches_dotted_label(group_by),
            ProjectionField::Scalar { .. } => false,
        })
        && aggregate_fields.iter().all(|field| {
            matches!(
                field,
                ProjectionField::Scalar {
                    expr: Expr::Aggregate(_),
                    alias: None,
                }
            )
        })
}

// Enforce the SQL DISTINCT rule at lowering time: every ORDER BY term must be
// fully expressible as a function of the outward projected distinct tuple
// rather than from hidden base-row fields.
pub(super) fn validate_distinct_order_terms_against_projection(
    projection: &ProjectionSelection,
    order_by: &[LoweredSqlOrderTerm],
    schema: &SchemaInfo,
) -> Result<(), SqlLoweringError> {
    if order_by
        .iter()
        .all(|term| distinct_order_term_is_derivable_from_projection(projection, term, schema))
    {
        return Ok(());
    }

    Err(SqlLoweringError::distinct_order_by_requires_projected_tuple())
}

// Keep DISTINCT ORDER BY derivation intentionally narrow:
// - exact projected expressions are always admissible
// - otherwise the term may reference only direct projected fields
// - `SELECT DISTINCT *` exposes the full entity field set
// This is a field-level proof only; it does not admit hidden payloads or a
// broader symbolic derivation model.
fn distinct_order_term_is_derivable_from_projection(
    projection: &ProjectionSelection,
    order_term: &LoweredSqlOrderTerm,
    schema: &SchemaInfo,
) -> bool {
    match projection {
        ProjectionSelection::All => {
            let projected_fields = schema.field_names_in_slot_order();

            order_term
                .analysis
                .references_only_direct_fields(projected_fields.as_slice())
        }
        ProjectionSelection::Fields(field_ids) => {
            let projected_fields = field_ids.iter().map(FieldId::as_str).collect::<Vec<_>>();

            order_term
                .analysis
                .references_only_direct_fields(projected_fields.as_slice())
        }
        ProjectionSelection::Exprs(fields) => {
            if fields.iter().any(|field| match field {
                ProjectionField::Scalar { expr, .. } => expr == &order_term.expr,
            }) {
                return true;
            }

            let projected_fields = fields
                .iter()
                .filter_map(ProjectionField::direct_field_name)
                .collect::<Vec<_>>();

            order_term
                .analysis
                .references_only_direct_fields(projected_fields.as_slice())
        }
    }
}

pub(super) fn direct_scalar_field_selection(
    items: &[SqlSelectItem],
    projection_aliases: &[Option<String>],
) -> Option<Vec<FieldId>> {
    if !projection_aliases.iter().all(Option::is_none) {
        return None;
    }

    items
        .iter()
        .map(|item| match item {
            SqlSelectItem::Field(field) if split_qualified_identifier(field).is_none() => {
                Some(FieldId::new(field.clone()))
            }
            SqlSelectItem::Field(_) | SqlSelectItem::Aggregate(_) | SqlSelectItem::Expr(_) => None,
        })
        .collect()
}

pub(in crate::db::sql::lowering) fn lower_select_item_expr(
    item: &SqlSelectItem,
    phase: SqlExprPhase,
    work: &PreparationWork<'_>,
) -> Result<Expr, SqlLoweringError> {
    lower_sql_expr(
        &crate::db::sql::parser::SqlExpr::from_select_item(item),
        phase,
        work,
    )
}

pub(in crate::db::sql::lowering) fn lower_analyzed_select_item_expr(
    item: &SqlSelectItem,
    phase: SqlExprPhase,
    work: &PreparationWork<'_>,
) -> Result<AnalyzedLoweredExpr, SqlLoweringError> {
    Ok(AnalyzedLoweredExpr::new(lower_select_item_expr(
        item, phase, work,
    )?))
}

#[cfg(test)]
mod tests {
    use crate::db::{
        query::plan::expr::{Alias, Expr, FieldPath, ProjectionField},
        sql::lowering::select::projection::grouped_projection_is_canonical_identity,
    };

    #[test]
    fn grouped_path_identity_preserves_exact_labels_and_alias_policy() {
        let mut fields = [ProjectionField::Scalar {
            expr: Expr::FieldPath(FieldPath::new("profile", vec!["city".into()])),
            alias: None,
        }];
        assert!(grouped_projection_is_canonical_identity(
            &fields,
            &["profile.city".into()],
        ));
        for label in ["profile", "profile.city.code", "profile.city."] {
            assert!(!grouped_projection_is_canonical_identity(
                &fields,
                &[label.into()],
            ));
        }
        let ProjectionField::Scalar { alias, .. } = &mut fields[0];
        *alias = Some(Alias::new("city"));
        assert!(!grouped_projection_is_canonical_identity(
            &fields,
            &["profile.city".into()],
        ));
    }
}