1use std::collections::HashMap;
41
42use rudb_common::{LogicalType, Result, Value};
43use rudb_kernels::{Comparison, Connective, call_values, cast_value, combine, compare_values};
44use rudb_plan::{CompareOp, ConjunctionOp, Expr, ExprRef, Node, NodeRef, Plan, Slice, SortKey};
45use rudb_vector::Vector;
46
47use crate::pass::{Context, Pass, top_down};
48use crate::walk;
49
50pub const VOLATILE: [&str; 17] = [
57 "current_connection_id",
58 "current_query",
59 "current_query_id",
60 "current_transaction_id",
61 "currval",
62 "error",
63 "gen_random_uuid",
64 "nextval",
65 "random",
66 "setseed",
67 "setval",
68 "sleep_ms",
69 "stats",
70 "uuid",
71 "uuidv4",
72 "uuidv7",
73 "write_log",
74];
75
76#[derive(Debug, Clone, Copy)]
78pub struct ExpressionRewriter;
79
80impl Pass for ExpressionRewriter {
81 fn name(&self) -> &'static str {
82 "expression_rewriter"
83 }
84
85 fn run(&self, plan: &mut Plan, _context: &Context) -> Result<()> {
86 rewrite(plan);
87 Ok(())
88 }
89}
90
91type Done = HashMap<ExprRef, ExprRef>;
99
100fn rewrite(plan: &mut Plan) {
102 let mut done = Done::new();
103 for node in top_down(plan) {
104 node_expressions(plan, node, &mut done);
105 }
106}
107
108fn node_expressions(plan: &mut Plan, node: NodeRef, done: &mut Done) {
115 match *plan.node(node) {
116 Node::Get { .. }
117 | Node::Dummy
118 | Node::Limit { .. }
119 | Node::SetOp { .. }
120 | Node::CrossProduct { .. } => {}
121 Node::Values { rows, .. } => {
122 let held = plan.row_list(rows).to_vec();
123 let rewritten: Vec<Slice> =
124 held.iter().map(|&row| expr_list(plan, row, done).unwrap_or(row)).collect();
125 if rewritten != held {
126 let rows = plan.add_rows(&rewritten);
127 match plan.node_mut(node) {
128 Node::Values { rows: held, .. } => *held = rows,
129 _ => unreachable!("the node was a values list a moment ago"),
130 }
131 }
132 }
133 Node::TableFunction { args, .. } => {
134 if let Some(rewritten) = expr_list(plan, args, done) {
135 match plan.node_mut(node) {
136 Node::TableFunction { args, .. } => *args = rewritten,
137 _ => unreachable!("the node was a table function a moment ago"),
138 }
139 }
140 }
141 Node::Filter { predicate, .. } => {
142 let rewritten = expression(plan, predicate, done);
143 if rewritten != predicate {
144 match plan.node_mut(node) {
145 Node::Filter { predicate, .. } => *predicate = rewritten,
146 _ => unreachable!("the node was a filter a moment ago"),
147 }
148 }
149 }
150 Node::Project { exprs, .. } => {
151 if let Some(rewritten) = expr_list(plan, exprs, done) {
152 match plan.node_mut(node) {
153 Node::Project { exprs, .. } => *exprs = rewritten,
154 _ => unreachable!("the node was a projection a moment ago"),
155 }
156 }
157 }
158 Node::Aggregate { groups, aggregates, .. } => {
159 let rewritten_groups = expr_list(plan, groups, done);
160 let rewritten_aggregates = expr_list(plan, aggregates, done);
161 match plan.node_mut(node) {
162 Node::Aggregate { groups, aggregates, .. } => {
163 if let Some(rewritten) = rewritten_groups {
164 *groups = rewritten;
165 }
166 if let Some(rewritten) = rewritten_aggregates {
167 *aggregates = rewritten;
168 }
169 }
170 _ => unreachable!("the node was an aggregate a moment ago"),
171 }
172 }
173 Node::Sort { keys, .. } | Node::TopN { keys, .. } => {
174 let held = plan.sort_key_list(keys).to_vec();
175 let rewritten: Vec<SortKey> = held
176 .iter()
177 .map(|key| SortKey { expr: expression(plan, key.expr, done), ..*key })
178 .collect();
179 if rewritten != held {
180 let keys = plan.add_sort_keys(&rewritten);
181 match plan.node_mut(node) {
182 Node::Sort { keys: held, .. } | Node::TopN { keys: held, .. } => *held = keys,
183 _ => unreachable!("the node was a sort a moment ago"),
184 }
185 }
186 }
187 Node::Distinct { on, .. } => {
188 if let Some(rewritten) = expr_list(plan, on, done) {
189 match plan.node_mut(node) {
190 Node::Distinct { on, .. } => *on = rewritten,
191 _ => unreachable!("the node was a distinct a moment ago"),
192 }
193 }
194 }
195 Node::Join { conditions, .. } => {
196 if let Some(rewritten) = expr_list(plan, conditions, done) {
197 match plan.node_mut(node) {
198 Node::Join { conditions, .. } => *conditions = rewritten,
199 _ => unreachable!("the node was a join a moment ago"),
200 }
201 }
202 }
203 }
204}
205
206fn expr_list(plan: &mut Plan, slice: Slice, done: &mut Done) -> Option<Slice> {
208 walk::list(plan, slice, &mut |plan, expr| expression(plan, expr, done))
209}
210
211fn expression(plan: &mut Plan, expr: ExprRef, done: &mut Done) -> ExprRef {
219 if let Some(&already) = done.get(&expr) {
220 return already;
221 }
222 let rebuilt = walk::rebuild(plan, expr, &mut |plan, child| expression(plan, child, done));
223 let simplified = simplify(plan, rebuilt);
224 done.insert(expr, simplified);
225 simplified
226}
227
228fn simplify(plan: &mut Plan, expr: ExprRef) -> ExprRef {
230 if matches!(*plan.expr(expr), Expr::Constant(_)) {
231 return expr;
232 }
233 if let Some(value) = fold(plan, expr) {
234 if let Some(folded) = constant_of(plan, expr, value) {
235 return folded;
236 }
237 }
238 if let Some(widened) = widened_negation(plan, expr) {
239 return widened;
240 }
241 match *plan.expr(expr) {
242 Expr::Conjunction { op, children } => conjunction(plan, expr, op, children),
243 Expr::Case { arms, otherwise } => case(plan, expr, arms, otherwise),
244 Expr::Compare { op, left, right } => null_comparison(plan, expr, op, left, right),
245 _ => expr,
246 }
247}
248
249fn fold(plan: &Plan, expr: ExprRef) -> Option<Value> {
256 match *plan.expr(expr) {
257 Expr::Cast { input, try_cast } => {
258 let inner = constant(plan, input)?;
259 cast_value(&inner, plan.expr_type(expr), try_cast).ok()
260 }
261 Expr::Compare { op, left, right } => {
262 let left = constant(plan, left)?;
263 let right = constant(plan, right)?;
264 compare_values(comparison(op), &left, &right).ok()
265 }
266 Expr::Conjunction { op, children } => {
267 let values = constants(plan, children)?;
268 let vectors: Vec<Vector> = values
269 .into_iter()
270 .map(|value| Vector::constant(LogicalType::Boolean, value, 1))
271 .collect();
272 Some(combine(connective(op), &vectors).ok()?.value_at(0))
273 }
274 Expr::Function { name, args } => {
275 let name = plan.string(name);
276 if VOLATILE.contains(&name) {
277 return None;
278 }
279 let values = constants(plan, args)?;
280 call_values(name, &values, plan.expr_type(expr), None).ok()
285 }
286 _ => None,
287 }
288}
289
290fn constant(plan: &Plan, expr: ExprRef) -> Option<Value> {
292 match *plan.expr(expr) {
293 Expr::Constant(value) => Some(plan.value(value).clone()),
294 _ => None,
295 }
296}
297
298fn constants(plan: &Plan, slice: Slice) -> Option<Vec<Value>> {
300 plan.expr_list(slice).iter().map(|&expr| constant(plan, expr)).collect()
301}
302
303fn constant_of(plan: &mut Plan, expr: ExprRef, value: Value) -> Option<ExprRef> {
309 let ty = plan.expr_type(expr).clone();
310 if !value.is_null() && value.logical_type() != ty {
311 return None;
312 }
313 let held = plan.add_value(value);
314 Some(plan.add_expr(Expr::Constant(held), ty))
315}
316
317fn widened_negation(plan: &mut Plan, expr: ExprRef) -> Option<ExprRef> {
330 let Expr::Function { name, args } = *plan.expr(expr) else { return None };
331 if plan.string(name) != "-" {
332 return None;
333 }
334 let &[only] = plan.expr_list(args) else { return None };
335 let value = constant(plan, only)?;
336 if *plan.expr_type(expr) != value.logical_type() {
337 return None;
338 }
339 let widened = match value {
340 Value::TinyInt(i8::MIN) => Value::SmallInt(128),
341 Value::SmallInt(i16::MIN) => Value::Integer(32_768),
342 Value::Integer(i32::MIN) => Value::BigInt(2_147_483_648),
343 Value::BigInt(i64::MIN) => Value::HugeInt(9_223_372_036_854_775_808),
344 _ => return None,
345 };
346 let ty = widened.logical_type();
347 let held = plan.add_value(widened);
348 Some(plan.add_expr(Expr::Constant(held), ty))
349}
350
351fn conjunction(plan: &mut Plan, expr: ExprRef, op: ConjunctionOp, children: Slice) -> ExprRef {
356 let (decides, drops) = match op {
358 ConjunctionOp::And => (false, true),
359 ConjunctionOp::Or => (true, false),
360 };
361 let held = plan.expr_list(children).to_vec();
362 let mut kept = Vec::with_capacity(held.len());
363 for child in held.iter().copied() {
364 match constant(plan, child).as_ref().and_then(Value::as_bool) {
365 Some(known) if known == decides => {
366 return constant_of(plan, expr, Value::Boolean(decides)).unwrap_or(expr);
367 }
368 Some(_) => {}
369 None => kept.push(child),
370 }
371 }
372 if kept.len() == held.len() {
373 return expr;
374 }
375 match kept.as_slice() {
376 [] => constant_of(plan, expr, Value::Boolean(drops)).unwrap_or(expr),
377 [only] => *only,
378 rest => {
379 let children = plan.add_expr_list(rest);
380 plan.add_expr(Expr::Conjunction { op, children }, LogicalType::Boolean)
381 }
382 }
383}
384
385enum Fires {
387 Always,
389 Never,
392 Maybe,
394}
395
396fn fires(plan: &Plan, when: ExprRef) -> Fires {
397 match constant(plan, when) {
398 Some(Value::Boolean(true)) => Fires::Always,
399 Some(Value::Boolean(false) | Value::Null) => Fires::Never,
400 Some(_) | None => Fires::Maybe,
403 }
404}
405
406fn case(plan: &mut Plan, expr: ExprRef, arms: Slice, otherwise: Option<ExprRef>) -> ExprRef {
408 let held = plan.arm_list(arms).to_vec();
409 let mut kept = Vec::with_capacity(held.len());
410 let mut result = otherwise;
411 let mut cut = false;
412 for arm in held.iter().copied() {
413 match fires(plan, arm.when) {
414 Fires::Never => {}
415 Fires::Always => {
416 result = Some(arm.then);
417 cut = true;
418 break;
419 }
420 Fires::Maybe => kept.push(arm),
421 }
422 }
423 if kept.len() == held.len() && !cut {
424 return expr;
425 }
426 if kept.is_empty() {
427 return match result {
431 Some(only) if plan.expr_type(only) == plan.expr_type(expr) => only,
432 Some(_) => expr,
433 None => constant_of(plan, expr, Value::Null).unwrap_or(expr),
434 };
435 }
436 let ty = plan.expr_type(expr).clone();
437 let arms = plan.add_arms(&kept);
438 plan.add_expr(Expr::Case { arms, otherwise: result }, ty)
439}
440
441fn null_comparison(
450 plan: &mut Plan,
451 expr: ExprRef,
452 op: CompareOp,
453 left: ExprRef,
454 right: ExprRef,
455) -> ExprRef {
456 if matches!(op, CompareOp::DistinctFrom | CompareOp::NotDistinctFrom) {
457 return expr;
458 }
459 let is_null = |side| constant(plan, side).is_some_and(|value| value.is_null());
460 if is_null(left) || is_null(right) {
461 constant_of(plan, expr, Value::Null).unwrap_or(expr)
462 } else {
463 expr
464 }
465}
466
467fn comparison(op: CompareOp) -> Comparison {
476 match op {
477 CompareOp::Equal => Comparison::Equal,
478 CompareOp::NotEqual => Comparison::NotEqual,
479 CompareOp::Less => Comparison::Less,
480 CompareOp::LessOrEqual => Comparison::LessOrEqual,
481 CompareOp::Greater => Comparison::Greater,
482 CompareOp::GreaterOrEqual => Comparison::GreaterOrEqual,
483 CompareOp::DistinctFrom => Comparison::DistinctFrom,
484 CompareOp::NotDistinctFrom => Comparison::NotDistinctFrom,
485 }
486}
487
488fn connective(op: ConjunctionOp) -> Connective {
490 match op {
491 ConjunctionOp::And => Connective::And,
492 ConjunctionOp::Or => Connective::Or,
493 }
494}
495
496#[cfg(test)]
497mod tests {
498 use super::{ExpressionRewriter, VOLATILE};
499 use crate::pass::{Context, Pass};
500 use rudb_plan::Plan;
501
502 fn folded(text: &str) -> String {
504 let mut plan =
505 Plan::parse(text).unwrap_or_else(|error| panic!("{text} did not parse: {error}"));
506 ExpressionRewriter
507 .run(&mut plan, &Context::new())
508 .unwrap_or_else(|error| panic!("{text} did not fold: {error}"));
509 plan.validate().unwrap_or_else(|error| panic!("{text} folded to a bad plan: {error}"));
510 plan.to_string()
511 }
512
513 const SCAN: &str = " Get memory.main.t AS t #0 [a::INTEGER, b::VARCHAR, c::BOOLEAN]\n";
514
515 #[test]
516 fn arithmetic_over_constants_becomes_the_number() {
517 let before = format!("Project #1 [\"+\"(2::INTEGER, 3::INTEGER)::INTEGER AS n]\n{SCAN}");
518 let after = format!("Project #1 [5::INTEGER AS n]\n{SCAN}");
519 assert_eq!(folded(&before), after);
520 }
521
522 #[test]
523 fn a_nest_of_constants_folds_all_the_way_up_in_one_walk() {
524 let before = format!(
527 "Project #1 [\"+\"(\"+\"(1::INTEGER, 2::INTEGER)::INTEGER, 3::INTEGER)::INTEGER AS n]\n{SCAN}"
528 );
529 let after = format!("Project #1 [6::INTEGER AS n]\n{SCAN}");
530 assert_eq!(folded(&before), after);
531 }
532
533 #[test]
535 fn negating_the_smallest_value_of_a_signed_type_widens_by_one_step() {
536 let cases = [
537 ("-128::TINYINT", "TINYINT", "128::SMALLINT"),
538 ("-32768::SMALLINT", "SMALLINT", "32768::INTEGER"),
539 ("-2147483648::INTEGER", "INTEGER", "2147483648::BIGINT"),
540 ("-9223372036854775808::BIGINT", "BIGINT", "9223372036854775808::HUGEINT"),
541 ];
542 for (argument, ty, expected) in cases {
543 let before = format!("Project #1 [\"-\"({argument})::{ty} AS n]\n{SCAN}");
544 let after = format!("Project #1 [{expected} AS n]\n{SCAN}");
545 assert_eq!(folded(&before), after, "{argument}");
546 }
547 }
548
549 #[test]
550 fn negating_anything_but_the_smallest_value_keeps_the_type_it_was_given() {
551 let before = format!("Project #1 [\"-\"(-127::TINYINT)::TINYINT AS n]\n{SCAN}");
552 let after = format!("Project #1 [127::TINYINT AS n]\n{SCAN}");
553 assert_eq!(folded(&before), after);
554 let smallest = "-170141183460469231731687303715884105728::HUGEINT";
557 let hugeint = format!("Project #1 [\"-\"({smallest})::HUGEINT AS n]\n{SCAN}");
558 assert_eq!(folded(&hugeint), hugeint);
559 let column = format!("Project #1 [\"-\"(#0.0::INTEGER)::INTEGER AS n]\n{SCAN}");
561 assert_eq!(folded(&column), column);
562 }
563
564 #[test]
565 fn a_call_with_a_column_in_it_is_left_alone() {
566 let text = format!("Project #1 [\"+\"(#0.0::INTEGER, 3::INTEGER)::INTEGER AS n]\n{SCAN}");
567 assert_eq!(folded(&text), text);
568 }
569
570 #[test]
571 fn a_cast_of_a_constant_folds_and_one_that_would_raise_does_not() {
572 let before = format!("Project #1 [CAST('1'::VARCHAR)::INTEGER AS n]\n{SCAN}");
573 let after = format!("Project #1 [1::INTEGER AS n]\n{SCAN}");
574 assert_eq!(folded(&before), after);
575 let raises = format!("Project #1 [CAST('abc'::VARCHAR)::INTEGER AS n]\n{SCAN}");
578 assert_eq!(folded(&raises), raises);
579 }
580
581 #[test]
582 fn a_comparison_of_constants_becomes_a_boolean() {
583 let before = format!("Filter (1::INTEGER < 2::INTEGER)::BOOLEAN\n{SCAN}");
584 let after = format!("Filter TRUE::BOOLEAN\n{SCAN}");
585 assert_eq!(folded(&before), after);
586 }
587
588 #[test]
589 fn a_comparison_against_a_null_is_null_and_the_other_side_goes_with_it() {
590 let before = format!("Filter (#0.0::INTEGER = NULL::INTEGER)::BOOLEAN\n{SCAN}");
591 let after = format!("Filter NULL::BOOLEAN\n{SCAN}");
592 assert_eq!(folded(&before), after);
593 }
594
595 #[test]
596 fn the_two_comparisons_that_have_an_answer_over_a_null_keep_it() {
597 let text =
600 format!("Filter (#0.0::INTEGER IS NOT DISTINCT FROM NULL::INTEGER)::BOOLEAN\n{SCAN}");
601 assert_eq!(folded(&text), text);
602 }
603
604 #[test]
605 fn a_true_drops_out_of_an_and_and_a_false_decides_it() {
606 let before = format!("Filter (TRUE::BOOLEAN AND #0.2::BOOLEAN)::BOOLEAN\n{SCAN}");
607 let after = format!("Filter #0.2::BOOLEAN\n{SCAN}");
608 assert_eq!(folded(&before), after);
609 let decided = format!("Filter (FALSE::BOOLEAN AND #0.2::BOOLEAN)::BOOLEAN\n{SCAN}");
610 let all = format!("Filter FALSE::BOOLEAN\n{SCAN}");
611 assert_eq!(folded(&decided), all);
612 }
613
614 #[test]
615 fn a_false_drops_out_of_an_or_and_a_true_decides_it() {
616 let before = format!("Filter (FALSE::BOOLEAN OR #0.2::BOOLEAN)::BOOLEAN\n{SCAN}");
617 let after = format!("Filter #0.2::BOOLEAN\n{SCAN}");
618 assert_eq!(folded(&before), after);
619 let decided = format!("Filter (TRUE::BOOLEAN OR #0.2::BOOLEAN)::BOOLEAN\n{SCAN}");
620 let all = format!("Filter TRUE::BOOLEAN\n{SCAN}");
621 assert_eq!(folded(&decided), all);
622 }
623
624 #[test]
625 fn a_null_operand_of_an_and_is_kept_because_it_is_neither_the_answer_nor_the_operand() {
626 let text = format!("Filter (NULL::BOOLEAN AND #0.2::BOOLEAN)::BOOLEAN\n{SCAN}");
627 assert_eq!(folded(&text), text);
628 }
629
630 #[test]
631 fn a_conjunction_of_constants_is_the_three_valued_answer() {
632 let before = format!("Filter (NULL::BOOLEAN AND FALSE::BOOLEAN)::BOOLEAN\n{SCAN}");
635 let after = format!("Filter FALSE::BOOLEAN\n{SCAN}");
636 assert_eq!(folded(&before), after);
637 let other = format!("Filter (NULL::BOOLEAN OR TRUE::BOOLEAN)::BOOLEAN\n{SCAN}");
638 let answer = format!("Filter TRUE::BOOLEAN\n{SCAN}");
639 assert_eq!(folded(&other), answer);
640 }
641
642 #[test]
643 fn a_long_conjunction_keeps_the_operands_that_are_not_decided() {
644 let before = format!(
645 "Filter (#0.2::BOOLEAN AND TRUE::BOOLEAN AND (#0.0::INTEGER > 1::INTEGER)::BOOLEAN)::BOOLEAN\n{SCAN}"
646 );
647 let after = format!(
648 "Filter (#0.2::BOOLEAN AND (#0.0::INTEGER > 1::INTEGER)::BOOLEAN)::BOOLEAN\n{SCAN}"
649 );
650 assert_eq!(folded(&before), after);
651 }
652
653 #[test]
654 fn an_arm_that_cannot_fire_is_dropped_and_a_null_condition_is_one_of_them() {
655 let before = format!(
656 "Project #1 [CASE WHEN FALSE::BOOLEAN THEN 1::INTEGER ELSE #0.0::INTEGER END::INTEGER AS n]\n{SCAN}"
657 );
658 let after = format!("Project #1 [#0.0::INTEGER AS n]\n{SCAN}");
659 assert_eq!(folded(&before), after);
660 let null = format!(
661 "Project #1 [CASE WHEN NULL::BOOLEAN THEN 1::INTEGER ELSE #0.0::INTEGER END::INTEGER AS n]\n{SCAN}"
662 );
663 assert_eq!(folded(&null), after);
664 }
665
666 #[test]
667 fn the_first_arm_that_always_fires_cuts_the_ones_after_it() {
668 let before = format!(
669 "Project #1 [CASE WHEN #0.2::BOOLEAN THEN 1::INTEGER WHEN TRUE::BOOLEAN THEN 2::INTEGER WHEN #0.2::BOOLEAN THEN 3::INTEGER ELSE 4::INTEGER END::INTEGER AS n]\n{SCAN}"
670 );
671 let after = format!(
672 "Project #1 [CASE WHEN #0.2::BOOLEAN THEN 1::INTEGER ELSE 2::INTEGER END::INTEGER AS n]\n{SCAN}"
673 );
674 assert_eq!(folded(&before), after);
675 }
676
677 #[test]
678 fn a_case_with_no_arm_left_and_no_else_is_null() {
679 let before = format!(
680 "Project #1 [CASE WHEN FALSE::BOOLEAN THEN 1::INTEGER END::INTEGER AS n]\n{SCAN}"
681 );
682 let after = format!("Project #1 [NULL::INTEGER AS n]\n{SCAN}");
683 assert_eq!(folded(&before), after);
684 }
685
686 #[test]
687 fn an_aggregate_keeps_its_place_and_its_arguments_are_folded_under_it() {
688 let before = format!(
691 "Aggregate #1 groups=[] aggregates=[sum(\"+\"(1::INTEGER, 2::INTEGER)::INTEGER)::HUGEINT]\n{SCAN}"
692 );
693 let after = format!("Aggregate #1 groups=[] aggregates=[sum(3::INTEGER)::HUGEINT]\n{SCAN}");
694 assert_eq!(folded(&before), after);
695 }
696
697 #[test]
698 fn a_sort_key_and_a_join_condition_are_folded_too() {
699 let before =
700 format!("Sort [\"+\"(1::INTEGER, 1::INTEGER)::INTEGER ASC NULLS LAST]\n{SCAN}");
701 let after = format!("Sort [2::INTEGER ASC NULLS LAST]\n{SCAN}");
702 assert_eq!(folded(&before), after);
703 }
704
705 #[test]
706 fn folding_twice_is_folding_once() {
707 let before = format!(
708 "Filter (TRUE::BOOLEAN AND (\"+\"(1::INTEGER, 1::INTEGER)::INTEGER > #0.0::INTEGER)::BOOLEAN)::BOOLEAN\n{SCAN}"
709 );
710 let once = folded(&before);
711 assert_eq!(folded(&once), once);
712 }
713
714 #[test]
715 fn a_volatile_call_is_not_folded_however_constant_its_arguments_are() {
716 assert!(VOLATILE.contains(&"random"));
719 assert!(VOLATILE.contains(&"nextval"));
720 let text = format!("Project #1 [random()::DOUBLE AS n]\n{SCAN}");
721 assert_eq!(folded(&text), text);
722 }
723}