toasty 0.11.0

use super::Simplify;
use toasty_core::stmt::{self, Expr, ResolvedRef, VisitMut};

impl Simplify<'_> {
    /// Heavyweight binary-op rewrites. Cheap canonicalization (constant
    /// folding, null propagation, boolean-constant simplification,
    /// literal-on-right swap) runs in `fold::expr_binary_op` before this
    /// is reached, so heavyweight rules see operands in canonical form
    /// (no `(Value, Value)`, no `(Value, _)` ahead of `(_, Value)`).
    ///
    /// App-level rewrites on eq/ne operands (`Reference::Model` →
    /// primary-key field, `BelongsTo` → foreign-key field) fire in the
    /// pre-lowering `lower::expr_eq_operand::RewriteEqOperand` pass, not
    /// here.
    pub(super) fn simplify_expr_binary_op(
        &mut self,
        op: stmt::BinaryOp,
        lhs: &mut stmt::Expr,
        rhs: &mut stmt::Expr,
    ) -> Option<stmt::Expr> {
        if (op.is_eq() || op.is_ne())
            && let (Some(lhs_len), Some(rhs_len)) = (lhs.record_len(), rhs.record_len())
            && lhs_len != rhs_len
        {
            return Some(op.is_ne().into());
        }

        let result = match (&mut *lhs, &mut *rhs) {
            // Self-comparison, e.g.,
            //
            //  - `x = x` → `true`
            //  - `x != x` → `false`
            //
            // Only applied for non-nullable field references.
            (Expr::Reference(lhs), Expr::Reference(rhs))
                if lhs == rhs && (op.is_eq() || op.is_ne()) =>
            {
                if lhs.is_field() {
                    let field = self.cx.resolve_expr_reference(lhs).as_field_unwrap();
                    if !field.nullable() {
                        return Some(op.is_eq().into());
                    }
                }
                None
            }
            // Tuple decomposition,
            //
            //  - `(a, b) = (x, y)` → `a = x and b = y`
            //  - `(a, b) != (x, y)` → `a != x or b != y`
            (Expr::Record(lhs_rec), Expr::Record(rhs_rec))
                if (op.is_eq() || op.is_ne()) && lhs_rec.len() == rhs_rec.len() =>
            {
                let comparisons: Vec<_> = std::mem::take(&mut lhs_rec.fields)
                    .into_iter()
                    .zip(std::mem::take(&mut rhs_rec.fields))
                    .map(|(l, r)| record_field_comparison(l, op, r))
                    .collect();

                let mut comparison = if op.is_eq() {
                    Expr::and_from_vec(comparisons)
                } else {
                    Expr::or_from_vec(comparisons)
                };
                self.visit_expr_mut(&mut comparison);
                Some(comparison)
            }
            // Tuple decomposition with a Value::Record on one side,
            //
            //  - `(a, b) = Value::Record([x, y])` → `a = x and b = y`
            //
            // This arises after match elimination produces `Record([col1, col2]) == Value::Record([1, "alice"])`.
            (Expr::Record(rec), Expr::Value(stmt::Value::Record(val_rec)))
            | (Expr::Value(stmt::Value::Record(val_rec)), Expr::Record(rec))
                if (op.is_eq() || op.is_ne()) && rec.len() == val_rec.len() =>
            {
                let comparisons: Vec<_> = std::mem::take(&mut rec.fields)
                    .into_iter()
                    .zip(std::mem::take(&mut val_rec.fields))
                    .map(|(expr, val)| record_field_comparison(expr, op, Expr::from(val)))
                    .collect();

                let mut comparison = if op.is_eq() {
                    Expr::and_from_vec(comparisons)
                } else {
                    Expr::or_from_vec(comparisons)
                };
                self.visit_expr_mut(&mut comparison);
                Some(comparison)
            }
            // Match elimination: distribute binary op into match arms as OR
            //
            //   Match(subj, [p1 => e1, p2 => e2]) <op> rhs
            //   → OR(subj == p1 AND e1 <op> rhs, subj == p2 AND e2 <op> rhs)
            //
            // Each arm is fully simplified inline. Arms that fold to false/null
            // are pruned. Comparison ops only — distributing arithmetic this
            // way would produce a malformed boolean from a non-boolean term.
            (Expr::Match(m), _) if !op.is_arithmetic() && m.subject.is_stable() => {
                let match_expr = lhs.take();
                let other = rhs.take();
                Some(self.eliminate_match_in_binary_op(op, match_expr, other, true))
            }
            (_, Expr::Match(m)) if !op.is_arithmetic() && m.subject.is_stable() => {
                let other = lhs.take();
                let match_expr = rhs.take();
                Some(self.eliminate_match_in_binary_op(op, match_expr, other, false))
            }
            // Decode-cast stripping on comparisons with a constant. A stored
            // column whose type differs from the model type decodes through
            // `cast(col, <model ty>)` (see `map_table_column_to_model`). When
            // such a cast surfaces in a comparison — e.g. out of an enum
            // decode `Match` arm, which only unfolds via match elimination
            // after lowering has run — move the conversion onto the constant
            // side so the driver compares the stored form directly. Other
            // cast shapes are handled during lowering
            // (`lower_expr_binary_op`); this rule fires only on the
            // post-lower column shape.
            (Expr::Cast(cast), Expr::Value(value)) | (Expr::Value(value), Expr::Cast(cast))
                if (op.is_eq() || op.is_ne()) && cast.from.is_none() && cast.expr.is_column() =>
            {
                self.strip_decode_cast_comparison(op, cast, value)
            }
            // Decode-cast stripping on column-versus-column comparisons, such
            // as one embedded relation compared to another after both sides
            // substitute to key columns. The eligibility check proves that
            // comparing the stored forms preserves decoded equality.
            (Expr::Cast(lhs_cast), Expr::Cast(rhs_cast))
                if self.can_strip_decode_cast_column_comparison(op, lhs_cast, rhs_cast) =>
            {
                self.strip_decode_cast_column_comparison(op, lhs_cast, rhs_cast)
            }
            // Self-comparison with projections, e.g.,
            //
            //  - `address.city = address.city` → `true`
            //  - `address.city != address.city` → `false`
            //
            // By this point, constant projections and record projections have been simplified.
            // What remains are projections with opaque bases (e.g., field references).
            // `lhs.base.is_stable()` keeps this sound: a projection through a
            // non-deterministic base would evaluate the base twice and could
            // yield different values each time.
            (Expr::Project(lhs), Expr::Project(rhs))
                if lhs == rhs && lhs.base.is_stable() && (op.is_eq() || op.is_ne()) =>
            {
                // TODO: Check if the projected value is nullable
                Some(Expr::from(op.is_eq()))
            }
            _ => None,
        };

        if result.is_some() {
            return result;
        }

        // Null propagation for derived VALUES columns.
        //
        // If either operand is a column reference into a derived VALUES
        // table where every row has NULL at that column position, the
        // binary op can never produce a non-null result.
        if self.is_always_null_derived_column(lhs) || self.is_always_null_derived_column(rhs) {
            return Some(Expr::null());
        }

        // Relation-path-comparison and IN-subquery lifting fire in the
        // pre-lowering `lower::lift_in_subquery::*` pass, not here.
        None
    }

    /// Rewrites `cast(col, T) <eq/ne> const` so the conversion happens on
    /// the constant instead of the column: converts `const` to `col`'s
    /// stored type once, here, and emits `col <eq/ne> const'` so the driver
    /// compares the bare column against the converted constant rather than
    /// casting the column on every row.
    ///
    /// The target type comes from the referenced column, not from
    /// `Capability::native_type_for(T)`: a `#[column(type = ...)]` override
    /// can store the value in a different type than the backend's default
    /// (e.g. a text-stored UUID on SQLite, whose default UUID storage is a
    /// blob). Bails on anything that does not resolve to a physical column.
    fn strip_decode_cast_comparison(
        &mut self,
        op: stmt::BinaryOp,
        cast: &mut stmt::ExprCast,
        value: &mut stmt::Value,
    ) -> Option<Expr> {
        let expr_reference = cast.expr.as_expr_reference()?;

        let ResolvedRef::Column(column) = self.cx.resolve_expr_reference(expr_reference) else {
            return None;
        };

        let value = column
            .ty
            .cast(self.cx.schema(), value.take())
            .expect("failed to cast value");
        Some(Expr::binary_op(cast.expr.take(), op, value))
    }

    /// Returns whether `cast(col_a, T) <eq/ne> cast(col_b, T)` can compare the
    /// stored columns directly without changing the result.
    fn can_strip_decode_cast_column_comparison(
        &self,
        op: stmt::BinaryOp,
        lhs: &stmt::ExprCast,
        rhs: &stmt::ExprCast,
    ) -> bool {
        if !(op.is_eq() || op.is_ne())
            || lhs.from.is_some()
            || rhs.from.is_some()
            || !lhs.expr.is_column()
            || !rhs.expr.is_column()
            || lhs.ty != rhs.ty
        {
            return false;
        }

        let Some(lhs_reference) = lhs.expr.as_expr_reference() else {
            return false;
        };
        let Some(rhs_reference) = rhs.expr.as_expr_reference() else {
            return false;
        };

        let ResolvedRef::Column(lhs_column) = self.cx.resolve_expr_reference(lhs_reference) else {
            return false;
        };
        let ResolvedRef::Column(rhs_column) = self.cx.resolve_expr_reference(rhs_reference) else {
            return false;
        };

        lhs_column.ty == rhs_column.ty && lhs_column.ty.cast_preserves_equality(&lhs.ty)
    }

    /// Rewrites `cast(col_a, T) <eq/ne> cast(col_b, T)` to
    /// `col_a <eq/ne> col_b` after
    /// [`Self::can_strip_decode_cast_column_comparison`] proves the rewrite is
    /// valid.
    fn strip_decode_cast_column_comparison(
        &mut self,
        op: stmt::BinaryOp,
        lhs: &mut stmt::ExprCast,
        rhs: &mut stmt::ExprCast,
    ) -> Option<Expr> {
        Some(Expr::binary_op(lhs.expr.take(), op, rhs.expr.take()))
    }

    /// Returns `true` if `expr` is a column reference that resolves to a
    /// derived VALUES table where every row has NULL at the referenced column.
    fn is_always_null_derived_column(&self, expr: &Expr) -> bool {
        let Expr::Reference(expr_ref) = expr else {
            return false;
        };

        match self.cx.resolve_expr_reference(expr_ref) {
            ResolvedRef::Derived(derived_ref) => derived_ref.is_column_always_null(),
            _ => false,
        }
    }

    /// Distributes a binary op over match arms, producing an OR of guarded
    /// comparisons. Each arm becomes `(subject == pattern) AND (arm_expr <op> other)`.
    /// Dead branches (false/null) are pruned after inline simplification.
    fn eliminate_match_in_binary_op(
        &mut self,
        op: stmt::BinaryOp,
        match_expr: Expr,
        other: Expr,
        match_on_lhs: bool,
    ) -> Expr {
        self.eliminate_match(match_expr, |arm| {
            if match_on_lhs {
                Expr::binary_op(arm, op, other.clone())
            } else {
                Expr::binary_op(other.clone(), op, arm)
            }
        })
    }

    /// Distributes a predicate over match arms, producing an OR of guarded
    /// terms. `term` builds the predicate for one arm from that arm's
    /// expression; each arm becomes `(subject == pattern) AND term(arm_expr)`,
    /// and the else branch is guarded by the negation of every pattern. Dead
    /// branches (false/null/error) are pruned after inline simplification.
    pub(super) fn eliminate_match(
        &mut self,
        match_expr: Expr,
        term: impl Fn(Expr) -> Expr,
    ) -> Expr {
        let Expr::Match(match_expr) = match_expr else {
            unreachable!()
        };

        let mut operands = Vec::new();

        // Collect arm patterns before consuming the arms (needed for the else guard).
        let patterns: Vec<_> = match_expr.arms.iter().map(|a| a.pattern.clone()).collect();

        for arm in match_expr.arms {
            // Unit enum variants return the discriminant subject. The arm's
            // pattern is the value of that subject within this branch.
            let arm_expr = if arm.expr == *match_expr.subject {
                Expr::from(arm.pattern.clone())
            } else {
                arm.expr
            };
            let guard = Expr::binary_op(
                (*match_expr.subject).clone(),
                stmt::BinaryOp::Eq,
                Expr::from(arm.pattern),
            );

            let mut term = Expr::and_from_vec(vec![guard, term(arm_expr)]);
            self.visit_expr_mut(&mut term);

            // Prune dead branches
            if is_dead_filter_term(&term) {
                continue;
            }

            operands.push(term);
        }

        // Include the else branch with a guard that negates all arm patterns.
        {
            let guards: Vec<Expr> = patterns
                .into_iter()
                .map(|pattern| {
                    Expr::not(Expr::binary_op(
                        (*match_expr.subject).clone(),
                        stmt::BinaryOp::Eq,
                        Expr::from(pattern),
                    ))
                })
                .collect();

            let mut else_operands = guards;
            else_operands.push(term(*match_expr.else_expr));
            let mut term = Expr::and_from_vec(else_operands);
            self.visit_expr_mut(&mut term);

            // Prune dead branches
            if !is_dead_filter_term(&term) {
                operands.push(term);
            }
        }

        Expr::or_from_vec(operands)
    }
}

/// Returns `true` when a distributed-comparison term can never be `true`, so it
/// adds nothing to the surrounding `OR` and can be dropped.
///
/// Three shapes qualify:
///
/// 1. The term is `false` or a bare `NULL`.
///
/// 2. The term is an `AND` with a `NULL` conjunct, such as `x AND NULL`. This
///    comes from comparing against a decode `Match`'s `null` else branch — for
///    example `eq(option_embed, Some(..))` on a `None` row. `x AND NULL` never
///    holds in three-valued logic, and DynamoDB rejects the bare `NULL`
///    placeholder it would otherwise serialize to.
///
/// 3. The term compares against an `Expr::Error`, the placeholder a total enum
///    decode puts in its unreachable else branch. For a two-variant enum,
///    distributing `field == "x"` over the decode produces a term like
///    `disc != 1 AND disc != 2 AND Error == "x"`. A surrounding variant gate
///    usually contradicts the `disc != k` guards and folds this away, but
///    factoring a shared predicate out from under its gates removes that gate
///    (issue #1061). Since `Error` marks a branch that never runs, the
///    comparison can never hold. Dropping the term also keeps `Error` out of
///    the SQL serializer, which cannot render it.
fn is_dead_filter_term(term: &Expr) -> bool {
    if term.is_unsatisfiable() {
        return true;
    }
    if contains_error(term) {
        return true;
    }
    matches!(
        term,
        Expr::And(and) if and.operands.iter().any(Expr::is_value_null)
    )
}

/// Returns `true` if `expr` contains an `Expr::Error` node anywhere in its
/// subtree — the unreachable-branch placeholder produced by enum decode.
fn contains_error(expr: &Expr) -> bool {
    use toasty_core::stmt::Visit;

    struct FindError(bool);

    impl Visit for FindError {
        fn visit_expr_error(&mut self, _: &stmt::ExprError) {
            self.0 = true;
        }
    }

    let mut find = FindError(false);
    find.visit_expr(expr);
    find.0
}

/// Uses null checks for record fields, including empty relation slots.
fn record_field_comparison(lhs: Expr, op: stmt::BinaryOp, rhs: Expr) -> Expr {
    let other = if lhs.is_value_null() {
        rhs
    } else if rhs.is_value_null() {
        lhs
    } else {
        return Expr::binary_op(lhs, op, rhs);
    };
    if op.is_eq() {
        Expr::is_null(other)
    } else {
        Expr::is_not_null(other)
    }
}