Skip to main content

drizzle_core/expr/
ops.rs

1//! Arithmetic operations using `std::ops` traits.
2//!
3//! This module implements `Add`, `Sub`, `Mul`, `Div`, `Rem` for `SQLExpr`,
4//! enabling natural Rust syntax for SQL arithmetic.
5
6use core::ops::{Add, Div, Mul, Neg, Rem, Sub};
7
8use crate::dialect::Dialect;
9use crate::sql::{SQL, SQLChunk, Token};
10use crate::traits::SQLParam;
11use crate::types::{AddOp, ArithmeticOutput, DivOp, MulOp, NegOutput, Numeric, RemOp, SubOp};
12
13use super::{AggOr, AggregateKind, Expr, Nullability, ResolveArithmeticNullability, SQLExpr};
14
15type ArithmeticNullable<'a, V, T, N, Rhs, Op> = <<T as ArithmeticOutput<
16    <Rhs as Expr<'a, V>>::SQLType,
17    Op,
18>>::Nullability as ResolveArithmeticNullability<
19    N,
20    <Rhs as Expr<'a, V>>::Nullable,
21>>::Output;
22
23#[inline]
24fn binary_op_sql<'a, V, L, R>(left: L, operator: Token, right: R) -> SQL<'a, V>
25where
26    V: SQLParam + 'a,
27    L: Expr<'a, V>,
28    R: Expr<'a, V>,
29{
30    binary_operator_sql(left.into_expr_sql(), operator, right.into_expr_sql())
31}
32
33// =============================================================================
34// Operator precedence
35// =============================================================================
36
37/// Rank given to anything between two operands that is not a ranked binary
38/// operator: comparisons, logical operators, raw operator text. It is below
39/// every ranked operator, so such an operand is always grouped.
40const LOOSEST: u8 = 0;
41
42/// How tightly `operator` binds between two operands in `dialect`; larger
43/// binds tighter. Only the operators the expression builders place between
44/// two operands are ranked.
45const fn binding_power(dialect: Dialect, operator: Token) -> u8 {
46    match dialect {
47        // SQLite ranks `||` above `*`, and `& | << >>` together below `+ -`.
48        Dialect::SQLite => match operator {
49            Token::CONCAT => 5,
50            Token::STAR | Token::SLASH | Token::REM => 4,
51            Token::PLUS | Token::MINUS => 3,
52            Token::BITAND | Token::BITOR | Token::LSHIFT | Token::RSHIFT => 2,
53            _ => LOOSEST,
54        },
55        // PostgreSQL puts `||` and the bitwise operators in its shared
56        // "any other operator" level, below `+ -`.
57        Dialect::PostgreSQL => match operator {
58            Token::STAR | Token::SLASH | Token::REM => 4,
59            Token::PLUS | Token::MINUS => 3,
60            Token::CONCAT | Token::BITAND | Token::BITOR | Token::LSHIFT | Token::RSHIFT => 2,
61            _ => LOOSEST,
62        },
63        // MySQL reads `||` as logical OR, so it stays unranked.
64        Dialect::MySQL => match operator {
65            Token::STAR | Token::SLASH | Token::REM => 6,
66            Token::PLUS | Token::MINUS => 5,
67            Token::LSHIFT | Token::RSHIFT => 4,
68            Token::BITAND => 3,
69            Token::BITOR => 2,
70            _ => LOOSEST,
71        },
72    }
73}
74
75/// The loosest-binding operator found at the top level of an operand.
76#[derive(Clone, Copy)]
77struct TopLevelOperator {
78    power: u8,
79    /// Every operator at `power` is `||`, which is associative.
80    only_concat: bool,
81}
82
83impl TopLevelOperator {
84    fn record(found: &mut Option<Self>, power: u8, concat: bool) {
85        match found {
86            None => {
87                *found = Some(Self {
88                    power,
89                    only_concat: concat,
90                });
91            }
92            Some(current) if power < current.power => {
93                *current = Self {
94                    power,
95                    only_concat: concat,
96                };
97            }
98            Some(current) if power == current.power => current.only_concat &= concat,
99            Some(_) => {}
100        }
101    }
102}
103
104/// Keyword and comparison tokens that join two operands at a looser level
105/// than any arithmetic operator.
106const fn is_loose_infix(token: Token) -> bool {
107    matches!(
108        token,
109        Token::EQ
110            | Token::NE
111            | Token::LT
112            | Token::GT
113            | Token::LE
114            | Token::GE
115            | Token::AND
116            | Token::OR
117            | Token::NOT
118            | Token::IS
119            | Token::ISNOT
120            | Token::IN
121            | Token::LIKE
122            | Token::BETWEEN
123            | Token::ESCAPE
124            | Token::ISNULL
125            | Token::NOTNULL
126            | Token::MATCH
127    )
128}
129
130/// Raw text that reads as a single term (a function name, keyword literal,
131/// or number) rather than as an operator.
132fn is_term_text(text: &str) -> bool {
133    let text = text.trim();
134    !text.is_empty()
135        && text
136            .chars()
137            .all(|ch| ch.is_alphanumeric() || matches!(ch, '_' | '.' | '"' | '`' | '\''))
138}
139
140/// Finds the loosest binary operator outside any parentheses or
141/// `CASE ... END` in `operand`. `None` means the operand is a single term: a
142/// column, value, function call, parenthesized group, or a sign-prefixed term.
143fn top_level_operator<V: SQLParam>(operand: &SQL<'_, V>) -> Option<TopLevelOperator> {
144    let mut depth = 0usize;
145    // True at the start and after an operator, where `-`/`+` are signs.
146    let mut expect_term = true;
147    let mut found = None;
148
149    for chunk in &operand.chunks {
150        match chunk {
151            SQLChunk::Token(Token::LPAREN | Token::CASE) => {
152                depth += 1;
153                expect_term = false;
154            }
155            SQLChunk::Token(Token::RPAREN | Token::END) => {
156                depth = depth.saturating_sub(1);
157                expect_term = false;
158            }
159            _ if depth > 0 => {}
160            SQLChunk::Token(
161                token @ (Token::PLUS
162                | Token::MINUS
163                | Token::STAR
164                | Token::SLASH
165                | Token::REM
166                | Token::CONCAT
167                | Token::BITAND
168                | Token::BITOR
169                | Token::LSHIFT
170                | Token::RSHIFT),
171            ) => {
172                if !expect_term {
173                    TopLevelOperator::record(
174                        &mut found,
175                        binding_power(V::DIALECT, *token),
176                        matches!(token, Token::CONCAT),
177                    );
178                    expect_term = true;
179                }
180            }
181            SQLChunk::Token(Token::BITNOT) => {}
182            SQLChunk::Token(token) if is_loose_infix(*token) => {
183                TopLevelOperator::record(&mut found, LOOSEST, false);
184                expect_term = true;
185            }
186            SQLChunk::Raw(text) if expect_term && matches!(text.trim(), "-" | "+") => {}
187            SQLChunk::Raw(text) if !is_term_text(text) => {
188                TopLevelOperator::record(&mut found, LOOSEST, false);
189                expect_term = true;
190            }
191            _ => expect_term = false,
192        }
193    }
194
195    found
196}
197
198/// Whether `operand` must be parenthesized to stay a single operand of
199/// `operator`. SQL operators of equal precedence associate to the left, so a
200/// right-hand operand also needs grouping at equal precedence (`a - (b - c)`
201/// is not `a - b - c`); the associative `||` is the exception.
202fn needs_grouping<V: SQLParam>(operand: &SQL<'_, V>, operator: Token, right_hand: bool) -> bool {
203    let Some(inner) = top_level_operator(operand) else {
204        return false;
205    };
206    let outer = binding_power(V::DIALECT, operator);
207    if right_hand {
208        inner.power < outer
209            || (inner.power == outer && !(matches!(operator, Token::CONCAT) && inner.only_concat))
210    } else {
211        inner.power < outer
212    }
213}
214
215/// Renders `left operator right` so the database evaluates the tree the Rust
216/// expression built.
217///
218/// An operand that is itself a binary expression is parenthesized when its
219/// top-level operator binds more loosely than `operator` (on the right-hand
220/// side, also when it binds equally), so `a * (b + c)` keeps its grouping.
221/// Operands that already read correctly stay flat: a single term renders as
222/// before, and so does a chain such as `a * b + c`.
223pub(crate) fn binary_operator_sql<'a, V>(
224    left: SQL<'a, V>,
225    operator: Token,
226    right: SQL<'a, V>,
227) -> SQL<'a, V>
228where
229    V: SQLParam + 'a,
230{
231    let left = left.parens_if_subquery();
232    let right = right.parens_if_subquery();
233    let left = if needs_grouping(&left, operator, false) {
234        left.parens()
235    } else {
236        left
237    };
238    let right = if needs_grouping(&right, operator, true) {
239        right.parens()
240    } else {
241        right
242    };
243    left.push(operator).append(right)
244}
245
246// =============================================================================
247// Addition
248// =============================================================================
249
250impl<'a, V, T, N, A, Rhs> Add<Rhs> for SQLExpr<'a, V, T, N, A>
251where
252    V: SQLParam + 'a,
253    T: ArithmeticOutput<Rhs::SQLType, AddOp>,
254    N: Nullability,
255    A: AggOr<Rhs::Aggregate>,
256    Rhs: Expr<'a, V>,
257    Rhs::SQLType: Numeric,
258    Rhs::Nullable: Nullability,
259    <T as ArithmeticOutput<Rhs::SQLType, AddOp>>::Nullability:
260        ResolveArithmeticNullability<N, Rhs::Nullable>,
261{
262    type Output = SQLExpr<
263        'a,
264        V,
265        <T as ArithmeticOutput<Rhs::SQLType, AddOp>>::Output,
266        ArithmeticNullable<'a, V, T, N, Rhs, AddOp>,
267        <A as AggOr<Rhs::Aggregate>>::Output,
268    >;
269
270    fn add(self, rhs: Rhs) -> Self::Output {
271        SQLExpr::new(binary_op_sql(self, Token::PLUS, rhs))
272    }
273}
274
275// =============================================================================
276// Subtraction
277// =============================================================================
278
279impl<'a, V, T, N, A, Rhs> Sub<Rhs> for SQLExpr<'a, V, T, N, A>
280where
281    V: SQLParam + 'a,
282    T: ArithmeticOutput<Rhs::SQLType, SubOp>,
283    N: Nullability,
284    A: AggOr<Rhs::Aggregate>,
285    Rhs: Expr<'a, V>,
286    Rhs::SQLType: Numeric,
287    Rhs::Nullable: Nullability,
288    <T as ArithmeticOutput<Rhs::SQLType, SubOp>>::Nullability:
289        ResolveArithmeticNullability<N, Rhs::Nullable>,
290{
291    type Output = SQLExpr<
292        'a,
293        V,
294        <T as ArithmeticOutput<Rhs::SQLType, SubOp>>::Output,
295        ArithmeticNullable<'a, V, T, N, Rhs, SubOp>,
296        <A as AggOr<Rhs::Aggregate>>::Output,
297    >;
298
299    fn sub(self, rhs: Rhs) -> Self::Output {
300        SQLExpr::new(binary_op_sql(self, Token::MINUS, rhs))
301    }
302}
303
304// =============================================================================
305// Multiplication
306// =============================================================================
307
308impl<'a, V, T, N, A, Rhs> Mul<Rhs> for SQLExpr<'a, V, T, N, A>
309where
310    V: SQLParam + 'a,
311    T: ArithmeticOutput<Rhs::SQLType, MulOp>,
312    N: Nullability,
313    A: AggOr<Rhs::Aggregate>,
314    Rhs: Expr<'a, V>,
315    Rhs::SQLType: Numeric,
316    Rhs::Nullable: Nullability,
317    <T as ArithmeticOutput<Rhs::SQLType, MulOp>>::Nullability:
318        ResolveArithmeticNullability<N, Rhs::Nullable>,
319{
320    type Output = SQLExpr<
321        'a,
322        V,
323        <T as ArithmeticOutput<Rhs::SQLType, MulOp>>::Output,
324        ArithmeticNullable<'a, V, T, N, Rhs, MulOp>,
325        <A as AggOr<Rhs::Aggregate>>::Output,
326    >;
327
328    fn mul(self, rhs: Rhs) -> Self::Output {
329        SQLExpr::new(binary_op_sql(self, Token::STAR, rhs))
330    }
331}
332
333// =============================================================================
334// Division
335// =============================================================================
336
337impl<'a, V, T, N, A, Rhs> Div<Rhs> for SQLExpr<'a, V, T, N, A>
338where
339    V: SQLParam + 'a,
340    T: ArithmeticOutput<Rhs::SQLType, DivOp>,
341    N: Nullability,
342    A: AggOr<Rhs::Aggregate>,
343    Rhs: Expr<'a, V>,
344    Rhs::SQLType: Numeric,
345    Rhs::Nullable: Nullability,
346    <T as ArithmeticOutput<Rhs::SQLType, DivOp>>::Nullability:
347        ResolveArithmeticNullability<N, Rhs::Nullable>,
348{
349    type Output = SQLExpr<
350        'a,
351        V,
352        <T as ArithmeticOutput<Rhs::SQLType, DivOp>>::Output,
353        ArithmeticNullable<'a, V, T, N, Rhs, DivOp>,
354        <A as AggOr<Rhs::Aggregate>>::Output,
355    >;
356
357    fn div(self, rhs: Rhs) -> Self::Output {
358        SQLExpr::new(binary_op_sql(self, Token::SLASH, rhs))
359    }
360}
361
362// =============================================================================
363// Remainder (Modulo)
364// =============================================================================
365
366impl<'a, V, T, N, A, Rhs> Rem<Rhs> for SQLExpr<'a, V, T, N, A>
367where
368    V: SQLParam + 'a,
369    T: ArithmeticOutput<Rhs::SQLType, RemOp>,
370    N: Nullability,
371    A: AggOr<Rhs::Aggregate>,
372    Rhs: Expr<'a, V>,
373    Rhs::SQLType: Numeric,
374    Rhs::Nullable: Nullability,
375    <T as ArithmeticOutput<Rhs::SQLType, RemOp>>::Nullability:
376        ResolveArithmeticNullability<N, Rhs::Nullable>,
377{
378    type Output = SQLExpr<
379        'a,
380        V,
381        <T as ArithmeticOutput<Rhs::SQLType, RemOp>>::Output,
382        ArithmeticNullable<'a, V, T, N, Rhs, RemOp>,
383        <A as AggOr<Rhs::Aggregate>>::Output,
384    >;
385
386    fn rem(self, rhs: Rhs) -> Self::Output {
387        SQLExpr::new(binary_op_sql(self, Token::REM, rhs))
388    }
389}
390
391// =============================================================================
392// Negation
393// =============================================================================
394
395impl<'a, V, T, N, A> Neg for SQLExpr<'a, V, T, N, A>
396where
397    V: SQLParam + 'a,
398    T: Numeric + NegOutput,
399    N: Nullability,
400    A: AggregateKind,
401{
402    type Output = SQLExpr<'a, V, T::Output, N, A>;
403
404    fn neg(self) -> Self::Output {
405        SQLExpr::new(SQL::from(Token::MINUS).append(self.into_expr_sql().parens()))
406    }
407}
408
409#[cfg(test)]
410mod tests {
411    use super::binary_operator_sql;
412    use crate::sql::{SQL, Token};
413    use crate::{Dialect, MySQLDialect, PostgresDialect, SQLParam, SQLiteDialect};
414
415    #[derive(Clone, Debug)]
416    struct SqliteParam;
417
418    impl SQLParam for SqliteParam {
419        const DIALECT: Dialect = Dialect::SQLite;
420        type DialectMarker = SQLiteDialect;
421    }
422
423    #[derive(Clone, Debug)]
424    struct PostgresParam;
425
426    impl SQLParam for PostgresParam {
427        const DIALECT: Dialect = Dialect::PostgreSQL;
428        type DialectMarker = PostgresDialect;
429    }
430
431    #[derive(Clone, Debug)]
432    struct MySqlParam;
433
434    impl SQLParam for MySqlParam {
435        const DIALECT: Dialect = Dialect::MySQL;
436        type DialectMarker = MySQLDialect;
437    }
438
439    fn term<V: SQLParam>(name: &'static str) -> SQL<'static, V> {
440        SQL::ident(name)
441    }
442
443    fn apply<V: SQLParam>(
444        left: SQL<'static, V>,
445        operator: Token,
446        right: SQL<'static, V>,
447    ) -> SQL<'static, V> {
448        binary_operator_sql(left, operator, right)
449    }
450
451    #[test]
452    fn single_operator_stays_flat() {
453        let product = apply::<SqliteParam>(term("a"), Token::STAR, term("b"));
454        assert_eq!(product.sql(), r#""a" * "b""#);
455    }
456
457    #[test]
458    fn looser_operand_is_grouped_on_either_side() {
459        let sum = || apply::<SqliteParam>(term("b"), Token::PLUS, term("c"));
460        assert_eq!(
461            apply(term("a"), Token::STAR, sum()).sql(),
462            r#""a" *("b" + "c")"#
463        );
464        assert_eq!(
465            apply(sum(), Token::STAR, term("a")).sql(),
466            r#"("b" + "c")* "a""#
467        );
468    }
469
470    #[test]
471    fn tighter_left_chain_stays_flat() {
472        let product = apply::<PostgresParam>(term("a"), Token::STAR, term("b"));
473        let chain = apply(product, Token::PLUS, term("c"));
474        assert_eq!(chain.sql(), r#""a" * "b" + "c""#);
475
476        let difference = apply::<PostgresParam>(term("a"), Token::MINUS, term("b"));
477        let chain = apply(difference, Token::MINUS, term("c"));
478        assert_eq!(chain.sql(), r#""a" - "b" - "c""#);
479    }
480
481    #[test]
482    fn equal_precedence_on_the_right_is_grouped() {
483        let difference = apply::<MySqlParam>(term("b"), Token::MINUS, term("c"));
484        assert_eq!(
485            apply(term("a"), Token::MINUS, difference).sql(),
486            "`a` -(`b` - `c`)"
487        );
488    }
489
490    #[test]
491    fn concatenation_chain_stays_flat_on_the_right() {
492        let tail = apply::<SqliteParam>(term("b"), Token::CONCAT, term("c"));
493        assert_eq!(
494            apply(term("a"), Token::CONCAT, tail).sql(),
495            r#""a" || "b" || "c""#
496        );
497    }
498
499    #[test]
500    fn concatenation_precedence_follows_the_dialect() {
501        // SQLite binds `||` tighter than `+`; PostgreSQL binds it looser.
502        let sqlite_sum = apply::<SqliteParam>(term("b"), Token::PLUS, term("c"));
503        assert_eq!(
504            apply(term("a"), Token::CONCAT, sqlite_sum).sql(),
505            r#""a" ||("b" + "c")"#
506        );
507
508        let postgres_sum = apply::<PostgresParam>(term("b"), Token::PLUS, term("c"));
509        assert_eq!(
510            apply(term("a"), Token::CONCAT, postgres_sum).sql(),
511            r#""a" || "b" + "c""#
512        );
513    }
514
515    #[test]
516    fn terms_with_inner_operators_stay_flat() {
517        // A function call, a parenthesized group and a signed term are each
518        // a single operand, whatever they contain.
519        let call = SQL::<SqliteParam>::raw("ABS")
520            .push(Token::LPAREN)
521            .append(apply(term("b"), Token::MINUS, term("c")))
522            .push(Token::RPAREN);
523        assert_eq!(
524            apply(term("a"), Token::STAR, call).sql(),
525            r#""a" * ABS ("b" - "c")"#
526        );
527
528        let signed = SQL::<SqliteParam>::raw("-").append(term("b"));
529        assert_eq!(
530            apply(term("a"), Token::MINUS, signed).sql(),
531            r#""a" - - "b""#
532        );
533    }
534
535    #[test]
536    fn comparison_operand_is_grouped() {
537        let comparison = SQL::<SqliteParam>::ident("b")
538            .push(Token::EQ)
539            .append(term("c"));
540        assert_eq!(
541            apply(term("a"), Token::PLUS, comparison).sql(),
542            r#""a" +("b" = "c")"#
543        );
544    }
545}