qraft-core 0.1.2

//! Binary predicates, comparisons, and numeric operators.

use std::marker::PhantomData;

use crate::{
    BinaryType, LowerCompatible, NullOf, Nullable, Numeric, Orderable, PredicateType,
    expression::{
        Column, Expression, PostfixOperator, Scalar,
        op::{And, Or},
    },
    instr::RpnInstr,
    lower::{Instructions, LowerCtx},
    ty::{Bool, Boolean, Comparable, Nullability, TypeMeta},
};

/// Operators emitted as binary instructions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Operator {
    Eq,
    Neq,
    Lt,
    Lte,
    Gt,
    Gte,
    Like { sensitive: bool, negated: bool },
    And,
    Or,
    Add,
    Mul,
    Rem,
    Sub,
    Div,
    BitAnd,
    BitOr,
    Shl,
    Shr,
}

/// A typed binary expression with a left and right operand.
pub struct Binary<T, L, R> {
    pub(crate) left: L,
    pub(crate) op: PhantomData<T>,
    pub(crate) right: R,
}

impl<T, L, R> Binary<T, L, R> {
    /// Returns the stored operands without changing their order.
    pub fn into_parts(self) -> (L, R) {
        (self.left, self.right)
    }
}

/// Maps operand types to the output type and emitted instruction.
pub trait BinaryOp<L: TypeMeta, R: TypeMeta> {
    /// Resulting SQL type for this operator and operand pair.
    type Output: TypeMeta;
    /// Lowered instruction template for this operator.
    const INSTRUCTION: RpnInstr;
}

/// Marker types for supported binary operators.
pub mod op {
    pub struct Eq;
    pub struct And;
    pub struct Or;
    pub struct Neq;
    pub struct Gt;
    pub struct Gte;
    pub struct Lt;
    pub struct Lte;
    pub struct Add;
    pub struct Rem;
    pub struct Mul;
    pub struct Sub;
    pub struct Div;
    pub struct BitAnd;
    pub struct BitOr;
    pub struct Shl;
    pub struct Shr;
}

macro_rules! impl_predicate_binary {
    ($op:ty, $variant:expr, $bound:ident) => {
        impl<T> BinaryOp<T, T> for $op
        where
            T: $bound + Nullability,
            BinaryType<NullOf<T>, NullOf<T>>: PredicateType,
        {
            type Output = <BinaryType<NullOf<T>, NullOf<T>> as PredicateType>::Output;
            const INSTRUCTION: RpnInstr = RpnInstr::Binary {
                op: $variant,
                lhs: 0,
                rhs: 0,
            };
        }
    };
}

impl_predicate_binary!(op::Eq, Operator::Eq, Comparable);
impl_predicate_binary!(op::Neq, Operator::Neq, Comparable);
impl_predicate_binary!(op::Gt, Operator::Gt, Orderable);
impl_predicate_binary!(op::Lt, Operator::Lt, Orderable);
impl_predicate_binary!(op::Gte, Operator::Gte, Orderable);
impl_predicate_binary!(op::Lte, Operator::Lte, Orderable);
impl_predicate_binary!(op::And, Operator::And, Boolean);
impl_predicate_binary!(op::Or, Operator::Or, Boolean);

macro_rules! impl_numeric_binary {
    ($op:ty, $variant:expr) => {
        impl<T> BinaryOp<T, T> for $op
        where
            T: Numeric,
        {
            type Output = T;
            const INSTRUCTION: RpnInstr = RpnInstr::Binary {
                op: $variant,
                lhs: 0,
                rhs: 0,
            };
        }
    };
}

impl_numeric_binary!(op::Add, Operator::Add);
impl_numeric_binary!(op::Rem, Operator::Rem);
impl_numeric_binary!(op::Mul, Operator::Mul);
impl_numeric_binary!(op::Sub, Operator::Sub);
impl_numeric_binary!(op::Div, Operator::Div);
impl_numeric_binary!(op::BitAnd, Operator::BitAnd);
impl_numeric_binary!(op::BitOr, Operator::BitOr);
impl_numeric_binary!(op::Shl, Operator::Shl);
impl_numeric_binary!(op::Shr, Operator::Shr);

#[qraft_expression_macro::as_expression]
impl<T, L, R> Expression for Binary<T, L, R>
where
    L: Expression,
    R: LowerCompatible<L::Type>,
    T: BinaryOp<L::Type, L::Type>,
{
    type Type = T::Output;

    fn lower(&self, ctx: &mut LowerCtx) -> usize {
        let lhs = self.left.lower(ctx);
        let rhs = self.right.lower_compatible(ctx);
        let RpnInstr::Binary { op, .. } = T::INSTRUCTION else {
            unreachable!("binary instruction");
        };
        ctx.instrs.push_binary(op, lhs, rhs);
        lhs + rhs + 1
    }
}

/// Equality and inequality helpers for comparable expressions.
pub trait EqExt<T: Comparable>: Sized + Expression {
    fn eq<E>(self, other: E) -> Binary<op::Eq, Self, E>
    where
        E: LowerCompatible<T>,
    {
        Binary {
            left: self,
            right: other,
            op: PhantomData,
        }
    }

    fn neq<E>(self, other: E) -> Binary<op::Neq, Self, E>
    where
        E: LowerCompatible<T>,
    {
        Binary {
            left: self,
            right: other,
            op: PhantomData,
        }
    }
}

/// Postfix predicates such as `is null` and `is true`.
#[derive(Debug)]
pub struct Postfix<L> {
    pub lhs: L,
    pub operator: PostfixOperator,
}

#[qraft_expression_macro::as_expression]
impl<E: Expression> Expression for Postfix<E> {
    type Type = Bool;

    fn lower(&self, ctx: &mut LowerCtx) -> usize {
        let lhs = self.lhs.lower(ctx);
        ctx.lower_postfix(self.operator, lhs)
    }
}

/// Adds `is null` and `is not null` to nullable expressions.
pub trait IsNull<T>: Sized {
    #[allow(clippy::wrong_self_convention)]
    fn is_null(self) -> Postfix<Self> {
        Postfix {
            lhs: self,
            operator: PostfixOperator::Null { negated: false },
        }
    }

    #[allow(clippy::wrong_self_convention)]
    fn is_not_null(self) -> Postfix<Self> {
        Postfix {
            lhs: self,
            operator: PostfixOperator::Null { negated: true },
        }
    }
}

/// Adds `is true` and `is false` to boolean expressions.
pub trait IsPredicate<T>: Sized {
    #[allow(clippy::wrong_self_convention)]
    fn is_true(self) -> Postfix<Self> {
        Postfix {
            lhs: self,
            operator: PostfixOperator::True,
        }
    }

    #[allow(clippy::wrong_self_convention)]
    fn is_false(self) -> Postfix<Self> {
        Postfix {
            lhs: self,
            operator: PostfixOperator::False,
        }
    }
}

impl<T: Comparable, V> IsNull<T> for V where V: Expression<Type = Nullable<T>> {}
impl<T: Boolean, V> IsPredicate<T> for V where V: Expression<Type = T> {}

impl<T: Comparable, V> EqExt<T> for V where V: Expression<Type = T> {}

/// Ordering predicates for comparable expressions.
pub trait OrderExt<T: Orderable>: Sized + Expression {
    fn lt<E>(self, other: E) -> Binary<op::Lt, Self, E>
    where
        E: LowerCompatible<T>,
    {
        Binary {
            left: self,
            right: other,
            op: PhantomData,
        }
    }

    fn lte<E>(self, other: E) -> Binary<op::Lte, Self, E>
    where
        E: LowerCompatible<T>,
    {
        Binary {
            left: self,
            right: other,
            op: PhantomData,
        }
    }

    fn gt<E>(self, other: E) -> Binary<op::Gt, Self, E>
    where
        E: LowerCompatible<T>,
    {
        Binary {
            left: self,
            right: other,
            op: PhantomData,
        }
    }

    fn gte<E>(self, other: E) -> Binary<op::Gte, Self, E>
    where
        E: LowerCompatible<T>,
    {
        Binary {
            left: self,
            right: other,
            op: PhantomData,
        }
    }
}

impl<T: Orderable, V> OrderExt<T> for V where V: Expression<Type = T> {}

/// Boolean combinators for predicate expressions.
pub trait PredicateExt<T: Boolean>: Sized + Expression {
    fn and<E>(self, other: E) -> Binary<And, Self, E>
    where
        E: LowerCompatible<T>,
    {
        Binary {
            left: self,
            right: other,
            op: PhantomData,
        }
    }

    fn or<E>(self, other: E) -> Binary<Or, Self, E>
    where
        E: LowerCompatible<T>,
    {
        Binary {
            left: self,
            right: other,
            op: PhantomData,
        }
    }
}

impl<T: Boolean, V> PredicateExt<T> for V where V: Expression<Type = T> {}

macro_rules! impl_binop_family {
    (trait = $trait:ident,method = $method:ident,op = $op:path) => {
        impl<T: Numeric, Op, L, R, E> core::ops::$trait<E> for Binary<Op, L, R>
        where
            E: Expression<Type = T>,
        {
            type Output = Binary<$op, Self, E>;

            fn $method(self, rhs: E) -> Self::Output {
                Binary {
                    left: self,
                    op: PhantomData,
                    right: rhs,
                }
            }
        }

        impl<T: Numeric, E> core::ops::$trait<E> for Scalar<T>
        where
            E: Expression<Type = T>,
        {
            type Output = Binary<$op, Self, E>;

            fn $method(self, rhs: E) -> Self::Output {
                Binary {
                    left: self,
                    op: PhantomData,
                    right: rhs,
                }
            }
        }

        impl<M, T: Numeric, E> core::ops::$trait<E> for Column<M, T>
        where
            E: Expression<Type = T>,
        {
            type Output = Binary<$op, Self, E>;

            fn $method(self, rhs: E) -> Self::Output {
                Binary {
                    left: self,
                    op: PhantomData,
                    right: rhs,
                }
            }
        }
    };
}

impl_binop_family!(trait = BitAnd, method = bitand, op = op::BitAnd);
impl_binop_family!(trait = BitOr, method = bitor, op = op::BitOr);
impl_binop_family!(trait = Shl, method = shl, op = op::Shl);
impl_binop_family!(trait = Shr, method = shr, op = op::Shr);