uqa-execution 0.5.3

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Projection-slot binding for supported scalar predicates.

use uqa_core::Value;
use uqa_sql::expr::cast_value;
use uqa_sql::{SQLError, SQLParam};

use super::{ProjectedExpr, ProjectedIntPredicate};
use crate::scalar::scalar_integer_binary_width;
use crate::{RowSchema, ScalarExpr};

#[expect(
    clippy::too_many_lines,
    reason = "predicate compiler exhaustively accepts or rejects each IR shape"
)]
pub(super) fn compile(
    expression: &ScalarExpr,
    schema: &RowSchema,
    params: &[SQLParam],
) -> Result<Option<ProjectedExpr>, SQLError> {
    let compiled = match expression {
        ScalarExpr::Column(column) => {
            let Some(index) = schema.unqualified_position(column) else {
                return Ok(None);
            };
            ProjectedExpr::Field(index)
        }
        ScalarExpr::Position(index) if *index < schema.len() => ProjectedExpr::Field(*index),
        ScalarExpr::QualifiedColumn { qualifier, column } => {
            let Some(index) = schema.qualified_position(qualifier, column) else {
                return Ok(None);
            };
            ProjectedExpr::Field(index)
        }
        ScalarExpr::Literal(value) | ScalarExpr::TypedLiteral { value, .. } => {
            ProjectedExpr::Literal(value.clone())
        }
        ScalarExpr::Param(index) => ProjectedExpr::Literal(parameter(*index, params)?),
        ScalarExpr::Binary { op, lhs, rhs } => {
            let integer_width = scalar_integer_binary_width(lhs, rhs, schema, params);
            let real_arithmetic = matches!(
                crate::scalar_type(expression, schema, params)?,
                Some(uqa_sql::ast::ColumnType::Real)
            );
            compiled_binary(
                *op,
                require(lhs, schema, params)?,
                require(rhs, schema, params)?,
                integer_width,
                real_arithmetic,
            )
        }
        ScalarExpr::UnaryMinus(expression) => {
            ProjectedExpr::UnaryMinus(Box::new(require(expression, schema, params)?))
        }
        ScalarExpr::Not(expression) => {
            ProjectedExpr::Not(Box::new(require(expression, schema, params)?))
        }
        ScalarExpr::And(items) => compiled_and(require_all(items, schema, params)?),
        ScalarExpr::Or(items) => ProjectedExpr::Or(require_all(items, schema, params)?),
        ScalarExpr::IsNull { expr, negated } => ProjectedExpr::IsNull {
            expression: Box::new(require(expr, schema, params)?),
            negated: *negated,
        },
        ScalarExpr::Between { expr, low, high } => compiled_between(
            require(expr, schema, params)?,
            require(low, schema, params)?,
            require(high, schema, params)?,
        ),
        ScalarExpr::InList {
            expr,
            list,
            negated,
        } => ProjectedExpr::InList {
            expression: Box::new(require(expr, schema, params)?),
            list: require_all(list, schema, params)?,
            negated: *negated,
        },
        ScalarExpr::Func {
            name,
            args,
            distinct,
            order_by,
            filter,
            ..
        } if !distinct && order_by.is_empty() && filter.is_none() => {
            let normalized = name.strip_prefix("pg_catalog.").unwrap_or(name);
            if !normalized.eq_ignore_ascii_case("like") && !normalized.eq_ignore_ascii_case("ilike")
            {
                return Ok(None);
            }
            let (expression, pattern, escape) = match args.as_slice() {
                [expression, pattern] => (expression, pattern, None),
                [expression, pattern, escape] => (expression, pattern, Some(escape)),
                _ => return Ok(None),
            };
            let pattern = match pattern {
                ScalarExpr::Literal(value) => value.clone(),
                ScalarExpr::Param(index) => parameter(*index, params)?,
                _ => return Ok(None),
            };
            if matches!(&pattern, Value::Null) {
                return Ok(Some(ProjectedExpr::Literal(Value::Null)));
            }
            let escape = match escape {
                Some(ScalarExpr::Literal(value)) => Some(value.clone()),
                Some(ScalarExpr::Param(index)) => Some(parameter(*index, params)?),
                Some(_) => return Ok(None),
                None => None,
            };
            if matches!(&escape, Some(Value::Null)) {
                return Ok(Some(ProjectedExpr::Literal(Value::Null)));
            }
            let pattern_text = uqa_sql::expr::value_to_string(&pattern)?;
            let escape_text = escape
                .as_ref()
                .map(uqa_sql::expr::value_to_string)
                .transpose()?;
            ProjectedExpr::Like {
                expression: Box::new(require(expression, schema, params)?),
                pattern: uqa_sql::expr::CompiledLikePattern::with_escape(
                    &pattern_text,
                    normalized.eq_ignore_ascii_case("ilike"),
                    escape_text.as_deref(),
                )?,
            }
        }
        ScalarExpr::Cast { expr, ty, .. } => {
            // A cast to a type whose input function reads the catalog, such as `regclass`, resolves names through the engine, which the projected form has no access to.
            if uqa_sql::ast::ColumnType::from_sql_name(ty)
                .map_or(true, |ty| uqa_sql::type_resolution::catalog_input_type(&ty))
            {
                return Ok(None);
            }
            let expression = require(expr, schema, params)?;
            match expression {
                // A typed SQL literal is represented as CAST(literal AS type).
                // Evaluate it once while preparing the predicate instead of
                // reparsing (notably DATE/TIMESTAMP text) for every input row.
                ProjectedExpr::Literal(value) if is_integer_type(ty) => ProjectedExpr::Cast {
                    expression: Box::new(ProjectedExpr::Literal(value)),
                    ty: ty.clone(),
                },
                ProjectedExpr::Literal(value) => ProjectedExpr::Literal(cast_value(&value, ty)?),
                expression => ProjectedExpr::Cast {
                    expression: Box::new(expression),
                    ty: ty.clone(),
                },
            }
        }
        ScalarExpr::Position(_)
        | ScalarExpr::InternalColumn(_)
        | ScalarExpr::Default
        | ScalarExpr::Star
        | ScalarExpr::QualifiedStar(_)
        | ScalarExpr::Func { .. }
        | ScalarExpr::Array(_)
        | ScalarExpr::Row(_)
        | ScalarExpr::CompositeRow { .. }
        | ScalarExpr::WindowCall { .. }
        | ScalarExpr::Case { .. }
        | ScalarExpr::ScalarSubquery(_)
        | ScalarExpr::Exists { .. }
        | ScalarExpr::InSubquery { .. } => return Ok(None),
    };
    Ok(Some(compiled))
}

fn compiled_and(items: Vec<ProjectedExpr>) -> ProjectedExpr {
    let all_integer_fields = items.iter().all(|item| {
        matches!(
            item,
            ProjectedExpr::IntFieldComparison { .. }
                | ProjectedExpr::IntFieldBetween { .. }
                | ProjectedExpr::IntFieldConjunction(_)
        )
    });
    if !all_integer_fields {
        return ProjectedExpr::And(items);
    }
    let mut predicates = Vec::with_capacity(items.len());
    for item in items {
        match item {
            ProjectedExpr::IntFieldComparison {
                field,
                op,
                literal,
                field_on_left,
            } => {
                push_integer_predicate(
                    &mut predicates,
                    ProjectedIntPredicate::Comparison {
                        field,
                        op,
                        literal,
                        field_on_left,
                    },
                );
            }
            ProjectedExpr::IntFieldBetween { field, low, high } => {
                push_integer_predicate(
                    &mut predicates,
                    ProjectedIntPredicate::Between { field, low, high },
                );
            }
            ProjectedExpr::IntFieldConjunction(items) => {
                for item in items {
                    push_integer_predicate(&mut predicates, item);
                }
            }
            _ => unreachable!("integer predicate conjunction was validated"),
        }
    }
    ProjectedExpr::IntFieldConjunction(predicates)
}

fn push_integer_predicate(
    predicates: &mut Vec<ProjectedIntPredicate>,
    next: ProjectedIntPredicate,
) {
    if let Some(range) = predicates
        .last()
        .and_then(|previous| adjacent_range(previous, &next))
    {
        *predicates
            .last_mut()
            .expect("matched the preceding comparison") = range;
    } else {
        predicates.push(next);
    }
}

/// Borrowing one field is stable within the input row. Coalesce only adjacent
/// inclusive lower/upper comparisons; no cast, effect or intervening predicate moves.
/// Commuted operands keep their original order in fallback error diagnostics.
fn adjacent_range(
    lower: &ProjectedIntPredicate,
    upper: &ProjectedIntPredicate,
) -> Option<ProjectedIntPredicate> {
    use uqa_sql::ast::BinaryOp::{GreaterEqual, LessEqual};
    let (
        ProjectedIntPredicate::Comparison {
            field: lower_field,
            op: GreaterEqual,
            literal: low,
            field_on_left: true,
        },
        ProjectedIntPredicate::Comparison {
            field: upper_field,
            op: LessEqual,
            literal: high,
            field_on_left: true,
        },
    ) = (lower, upper)
    else {
        return None;
    };
    (*lower_field == *upper_field).then_some(ProjectedIntPredicate::Between {
        field: *lower_field,
        low: *low,
        high: *high,
    })
}

fn is_integer_type(ty: &str) -> bool {
    matches!(
        ty,
        "smallint"
            | "int2"
            | "pg_catalog.int2"
            | "integer"
            | "int"
            | "int4"
            | "serial"
            | "serial4"
            | "pg_catalog.int4"
            | "bigint"
            | "int8"
            | "bigserial"
            | "serial8"
            | "pg_catalog.int8"
            | "oid"
            | "pg_catalog.oid"
            | "xid"
            | "pg_catalog.xid"
    )
}

fn compiled_binary(
    op: uqa_sql::ast::BinaryOp,
    lhs: ProjectedExpr,
    rhs: ProjectedExpr,
    integer_width: Option<uqa_sql::expr::IntegerWidth>,
    real_arithmetic: bool,
) -> ProjectedExpr {
    use uqa_sql::ast::BinaryOp;

    if matches!(
        op,
        BinaryOp::Equal
            | BinaryOp::NotEqual
            | BinaryOp::Less
            | BinaryOp::LessEqual
            | BinaryOp::Greater
            | BinaryOp::GreaterEqual
    ) {
        match (&lhs, &rhs) {
            (ProjectedExpr::Field(field), ProjectedExpr::Literal(Value::Int(literal))) => {
                return ProjectedExpr::IntFieldComparison {
                    field: *field,
                    op,
                    literal: *literal,
                    field_on_left: true,
                };
            }
            (ProjectedExpr::Literal(Value::Int(literal)), ProjectedExpr::Field(field)) => {
                return ProjectedExpr::IntFieldComparison {
                    field: *field,
                    op,
                    literal: *literal,
                    field_on_left: false,
                };
            }
            _ => {}
        }
    }
    ProjectedExpr::Binary {
        op,
        lhs: Box::new(lhs),
        rhs: Box::new(rhs),
        integer_width,
        real_arithmetic,
    }
}

fn compiled_between(
    expression: ProjectedExpr,
    low: ProjectedExpr,
    high: ProjectedExpr,
) -> ProjectedExpr {
    if let (
        ProjectedExpr::Field(field),
        ProjectedExpr::Literal(Value::Int(low)),
        ProjectedExpr::Literal(Value::Int(high)),
    ) = (&expression, &low, &high)
    {
        return ProjectedExpr::IntFieldBetween {
            field: *field,
            low: *low,
            high: *high,
        };
    }
    ProjectedExpr::Between {
        expression: Box::new(expression),
        low: Box::new(low),
        high: Box::new(high),
    }
}

fn require(
    expression: &ScalarExpr,
    schema: &RowSchema,
    params: &[SQLParam],
) -> Result<ProjectedExpr, SQLError> {
    compile(expression, schema, params)?.ok_or_else(|| {
        SQLError::Unsupported("expression cannot use positional predicate evaluation".into())
    })
}

fn require_all(
    expressions: &[ScalarExpr],
    schema: &RowSchema,
    params: &[SQLParam],
) -> Result<Vec<ProjectedExpr>, SQLError> {
    expressions
        .iter()
        .map(|expression| require(expression, schema, params))
        .collect()
}

fn parameter(index: usize, params: &[SQLParam]) -> Result<Value, SQLError> {
    match index.checked_sub(1).and_then(|offset| params.get(offset)) {
        Some(SQLParam::Scalar(value) | SQLParam::TypedScalar { value, .. }) => Ok(value.clone()),
        Some(SQLParam::Vector(_) | SQLParam::Tensor(_) | SQLParam::DeferredError { .. }) => Err(
            SQLError::Unsupported("parameter requires canonical predicate evaluation".into()),
        ),
        None => Err(SQLError::MissingParam(index)),
    }
}