core/optimizer: do expression rewriting on all expressions

Author: jussisaurio

core/translate/optimizer.rs

@@ -24,7 +24,7 @@ pub fn optimize_plan(plan: &mut Plan) -> Result<()> {
  */
 fn optimize_select_plan(plan: &mut SelectPlan) -> Result<()> {
     optimize_subqueries(&mut plan.source)?;
-    rewrite_exprs(&mut plan.source, &mut plan.where_clause)?;
+    rewrite_exprs_select(plan)?;
     if let ConstantConditionEliminationResult::ImpossibleCondition =
         eliminate_constants(&mut plan.source, &mut plan.where_clause)?
     {
@@ -55,7 +55,7 @@ fn optimize_select_plan(plan: &mut SelectPlan) -> Result<()> {
 }
 
 fn optimize_delete_plan(plan: &mut DeletePlan) -> Result<()> {
-    rewrite_exprs(&mut plan.source, &mut plan.where_clause)?;
+    rewrite_exprs_delete(plan)?;
     if let ConstantConditionEliminationResult::ImpossibleCondition =
         eliminate_constants(&mut plan.source, &mut plan.where_clause)?
     {
@@ -603,52 +603,83 @@ fn push_scan_direction(operator: &mut SourceOperator, direction: &Direction) {
     }
 }
 
-fn rewrite_exprs(
-    operator: &mut SourceOperator,
-    where_clauses: &mut Option<Vec<ast::Expr>>,
-) -> Result<()> {
-    if let Some(predicates) = where_clauses {
-        for expr in predicates.iter_mut() {
+fn rewrite_exprs_select(plan: &mut SelectPlan) -> Result<()> {
+    rewrite_source_operator_exprs(&mut plan.source)?;
+    for rc in plan.result_columns.iter_mut() {
+        rewrite_expr(&mut rc.expr)?;
+    }
+    for agg in plan.aggregates.iter_mut() {
+        rewrite_expr(&mut agg.original_expr)?;
+    }
+    if let Some(predicates) = &mut plan.where_clause {
+        for expr in predicates {
+            rewrite_expr(expr)?;
+        }
+    }
+    if let Some(group_by) = &mut plan.group_by {
+        for expr in group_by.exprs.iter_mut() {
+            rewrite_expr(expr)?;
+        }
+    }
+    if let Some(order_by) = &mut plan.order_by {
+        for (expr, _) in order_by.iter_mut() {
+            rewrite_expr(expr)?;
+        }
+    }
+
+    Ok(())
+}
+
+fn rewrite_exprs_delete(plan: &mut DeletePlan) -> Result<()> {
+    rewrite_source_operator_exprs(&mut plan.source)?;
+    if let Some(predicates) = &mut plan.where_clause {
+        for expr in predicates {
             rewrite_expr(expr)?;
         }
     }
 
+    Ok(())
+}
+
+fn rewrite_source_operator_exprs(operator: &mut SourceOperator) -> Result<()> {
     match operator {
         SourceOperator::Join {
             left,
             right,
             predicates,
             ..
         } => {
-            rewrite_exprs(left, where_clauses)?;
-            rewrite_exprs(right, where_clauses)?;
+            rewrite_source_operator_exprs(left)?;
+            rewrite_source_operator_exprs(right)?;
 
             if let Some(predicates) = predicates {
                 for expr in predicates.iter_mut() {
                     rewrite_expr(expr)?;
                 }
             }
+
+            Ok(())
         }
-        SourceOperator::Scan {
-            predicates: Some(preds),
-            ..
-        } => {
-            for expr in preds.iter_mut() {
-                rewrite_expr(expr)?;
+        SourceOperator::Scan { predicates, .. } | SourceOperator::Search { predicates, .. } => {
+            if let Some(predicates) = predicates {
+                for expr in predicates.iter_mut() {
+                    rewrite_expr(expr)?;
+                }
             }
+
+            Ok(())
         }
-        SourceOperator::Search {
-            predicates: Some(preds),
-            ..
-        } => {
-            for expr in preds.iter_mut() {
-                rewrite_expr(expr)?;
+        SourceOperator::Subquery { predicates, .. } => {
+            if let Some(predicates) = predicates {
+                for expr in predicates.iter_mut() {
+                    rewrite_expr(expr)?;
+                }
             }
+
+            Ok(())
         }
-        _ => (),
+        SourceOperator::Nothing { .. } => Ok(()),
     }
-
-    Ok(())
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]

core/translate/planner.rs

@@ -544,6 +544,14 @@ fn parse_join(
 pub fn parse_limit(limit: Limit) -> Option<usize> {
     if let Expr::Literal(ast::Literal::Numeric(n)) = limit.expr {
         n.parse().ok()
+    } else if let Expr::Id(id) = limit.expr {
+        if id.0.eq_ignore_ascii_case("true") {
+            Some(1)
+        } else if id.0.eq_ignore_ascii_case("false") {
+            Some(0)
+        } else {
+            None
+        }
     } else {
         None
     }

testing/select.test

@@ -11,6 +11,14 @@ do_execsql_test select-const-2 {
   SELECT 2
 } {2}
 
+do_execsql_test select-true {
+  SELECT true
+} {1}
+
+do_execsql_test select-false {
+  SELECT false
+} {0}
+
 do_execsql_test select-text-escape-1 {
   SELECT '''a'
 } {'a}
@@ -31,6 +39,15 @@ do_execsql_test select-limit-0 {
   SELECT id FROM users LIMIT 0;
 } {}
 
+# ORDER BY id here because sqlite uses age_idx here and we (yet) don't so force it to evaluate in ID order
+do_execsql_test select-limit-true {
+  SELECT id FROM users ORDER BY id LIMIT true;
+} {1}
+
+do_execsql_test select-limit-false {
+  SELECT id FROM users ORDER BY id LIMIT false;
+} {}
+
 do_execsql_test realify {
     select price from products limit 1;
 } {79.0}