Refactor planner and optimizer to be DRY

core/lib.rs

@@ -34,12 +34,12 @@ pub use storage::wal::WalFile;
 pub use storage::wal::WalFileShared;
 use util::parse_schema_rows;
 
-use translate::optimizer::optimize_plan;
 use translate::planner::prepare_select_plan;
 
 pub use error::LimboError;
 pub type Result<T> = std::result::Result<T, error::LimboError>;
 
+use crate::translate::optimizer::optimize_select_plan;
 pub use io::OpenFlags;
 #[cfg(feature = "fs")]
 pub use io::PlatformIO;
@@ -267,7 +267,7 @@ impl Connection {
                     match stmt {
                         ast::Stmt::Select(select) => {
                             let plan = prepare_select_plan(&*self.schema.borrow(), select)?;
-                            let plan = optimize_plan(plan)?;
+                            let plan = optimize_select_plan(plan)?;
                             println!("{}", plan);
                         }
                         _ => todo!(),

core/translate/optimizer.rs

@@ -9,45 +9,58 @@ use super::plan::{
     Direction, IterationDirection, Plan, Search, SourceOperator,
 };
 
+pub fn optimize_select_plan(plan: Plan) -> Result<Plan> {
+    optimize_plan(plan, true, true, true)
+}
+
+pub fn optimize_delete_plan(plan: Plan) -> Result<Plan> {
+    optimize_plan(plan, false, true, false)
+}
+
 /**
  * Make a few passes over the plan to optimize it.
  * TODO: these could probably be done in less passes,
  * but having them separate makes them easier to understand
  */
-pub fn optimize_plan(mut select_plan: Plan) -> Result<Plan> {
+fn optimize_plan(
+    mut select_plan: Plan,
+    optimize_push_predicates: bool,
+    optimize_use_indexes: bool,
+    optimize_eliminate_unnecessary_order_by: bool,
+) -> Result<Plan> {
     eliminate_between(&mut select_plan.source, &mut select_plan.where_clause)?;
     if let ConstantConditionEliminationResult::ImpossibleCondition =
         eliminate_constants(&mut select_plan.source, &mut select_plan.where_clause)?
     {
         select_plan.contains_constant_false_condition = true;
         return Ok(select_plan);
     }
-    push_predicates(
-        &mut select_plan.source,
-        &mut select_plan.where_clause,
-        &select_plan.referenced_tables,
-    )?;
-    use_indexes(
-        &mut select_plan.source,
-        &select_plan.referenced_tables,
-        &select_plan.available_indexes,
-    )?;
-    eliminate_unnecessary_orderby(
-        &mut select_plan.source,
-        &mut select_plan.order_by,
-        &select_plan.referenced_tables,
-        &select_plan.available_indexes,
-    )?;
-    Ok(select_plan)
-}
 
-pub fn optimize_delete_plan(mut delete_plan: Plan) -> Result<Plan> {
-    use_indexes(
-        &mut delete_plan.source,
-        &delete_plan.referenced_tables,
-        &delete_plan.available_indexes,
-    )?;
-    Ok(delete_plan)
+    if optimize_push_predicates {
+        push_predicates(
+            &mut select_plan.source,
+            &mut select_plan.where_clause,
+            &select_plan.referenced_tables,
+        )?;
+    }
+
+    if optimize_use_indexes {
+        use_indexes(
+            &mut select_plan.source,
+            &select_plan.referenced_tables,
+            &select_plan.available_indexes,
+        )?;
+    }
+
+    if optimize_eliminate_unnecessary_order_by {
+        eliminate_unnecessary_orderby(
+            &mut select_plan.source,
+            &mut select_plan.order_by,
+            &select_plan.referenced_tables,
+            &select_plan.available_indexes,
+        )?;
+    }
+    Ok(select_plan)
 }
 
 fn _operator_is_already_ordered_by(

core/translate/planner.rs

@@ -283,14 +283,7 @@ pub fn prepare_select_plan<'a>(schema: &Schema, select: ast::Select) -> Result<P
             };
 
             // Parse the WHERE clause
-            if let Some(w) = where_clause {
-                let mut predicates = vec![];
-                break_predicate_at_and_boundaries(w, &mut predicates);
-                for expr in predicates.iter_mut() {
-                    bind_column_references(expr, &plan.referenced_tables)?;
-                }
-                plan.where_clause = Some(predicates);
-            }
+            plan.where_clause = parse_where(where_clause, &plan.referenced_tables)?;
 
             let mut aggregate_expressions = Vec::new();
             for column in columns.clone() {
@@ -491,6 +484,49 @@ pub fn prepare_select_plan<'a>(schema: &Schema, select: ast::Select) -> Result<P
     }
 }
 
+pub fn prepare_delete_plan(
+    schema: &Schema,
+    tbl_name: &QualifiedName,
+    where_clause: Option<Expr>,
+) -> Result<Plan> {
+    let table_name = tbl_name.name.0.clone();
+
+    let table = if let Some(table) = schema.get_table(&table_name) {
+        table
+    } else {
+        crate::bail_parse_error!("Table {} not found", table_name);
+    };
+
+    let table_ref = BTreeTableReference {
+        table: table.clone(),
+        table_identifier: table_name.clone(),
+        table_index: 0,
+    };
+
+    // Parse and resolve the where_clause
+    let resolved_where_clauses = parse_where(where_clause, &[table_ref.clone()])?;
+
+    let plan = Plan {
+        source: SourceOperator::Scan {
+            id: 0,
+            table_reference: table_ref.clone(),
+            predicates: resolved_where_clauses.clone(),
+            iter_dir: None,
+        },
+        result_columns: vec![],
+        where_clause: resolved_where_clauses,
+        group_by: None,
+        order_by: None,
+        aggregates: vec![],
+        limit: None, // TODO: add support for limit
+        referenced_tables: vec![table_ref],
+        available_indexes: vec![],
+        contains_constant_false_condition: false,
+    };
+
+    Ok(plan)
+}
+
 #[allow(clippy::type_complexity)]
 fn parse_from(
     schema: &Schema,
@@ -552,6 +588,22 @@ fn parse_from(
     Ok((operator, tables))
 }
 
+fn parse_where(
+    where_clause: Option<Expr>,
+    referenced_tables: &[BTreeTableReference],
+) -> Result<Option<Vec<Expr>>> {
+    if let Some(where_expr) = where_clause {
+        let mut predicates = vec![];
+        break_predicate_at_and_boundaries(where_expr, &mut predicates);
+        for expr in predicates.iter_mut() {
+            bind_column_references(expr, referenced_tables)?;
+        }
+        Ok(Some(predicates))
+    } else {
+        Ok(None)
+    }
+}
+
 fn parse_join(
     schema: &Schema,
     join: ast::JoinedSelectTable,
@@ -735,57 +787,6 @@ fn parse_join(
     ))
 }
 
-pub fn prepare_delete_plan(
-    schema: &Schema,
-    tbl_name: &QualifiedName,
-    where_clause: Option<Expr>,
-) -> Result<Plan> {
-    let table_name = tbl_name.name.0.clone();
-
-    let table = if let Some(table) = schema.get_table(&table_name) {
-        table
-    } else {
-        crate::bail_parse_error!("Table {} not found", table_name);
-    };
-
-    let table_ref = BTreeTableReference {
-        table: table.clone(),
-        table_identifier: table_name.clone(),
-        table_index: 0,
-    };
-
-    // Parse and resolve the where_clause
-    let mut resolved_where_clause = None;
-    if let Some(where_expr) = where_clause {
-        let mut predicates = vec![];
-        break_predicate_at_and_boundaries(where_expr, &mut predicates);
-        for expr in predicates.iter_mut() {
-            bind_column_references(expr, &[table_ref.clone()])?;
-        }
-        resolved_where_clause = Some(predicates);
-    }
-
-    let plan = Plan {
-        source: SourceOperator::Scan {
-            id: 0,
-            table_reference: table_ref.clone(),
-            predicates: resolved_where_clause.clone(),
-            iter_dir: None,
-        },
-        result_columns: vec![],
-        where_clause: resolved_where_clause,
-        group_by: None,
-        order_by: None,
-        aggregates: vec![],
-        limit: None,
-        referenced_tables: vec![table_ref],
-        available_indexes: vec![],
-        contains_constant_false_condition: false,
-    };
-
-    Ok(plan)
-}
-
 fn break_predicate_at_and_boundaries(predicate: ast::Expr, out_predicates: &mut Vec<ast::Expr>) {
     match predicate {
         ast::Expr::Binary(left, ast::Operator::And, right) => {

core/translate/select.rs

@@ -1,22 +1,21 @@
 use std::rc::Weak;
 use std::{cell::RefCell, rc::Rc};
 
+use super::emitter::emit_program;
+use super::planner::prepare_select_plan;
 use crate::storage::sqlite3_ondisk::DatabaseHeader;
+use crate::translate::optimizer::optimize_select_plan;
 use crate::Connection;
 use crate::{schema::Schema, vdbe::Program, Result};
 use sqlite3_parser::ast;
 
-use super::emitter::emit_program;
-use super::optimizer::optimize_plan;
-use super::planner::prepare_select_plan;
-
 pub fn translate_select(
     schema: &Schema,
     select: ast::Select,
     database_header: Rc<RefCell<DatabaseHeader>>,
     connection: Weak<Connection>,
 ) -> Result<Program> {
     let select_plan = prepare_select_plan(schema, select)?;
-    let optimized_plan = optimize_plan(select_plan)?;
+    let optimized_plan = optimize_select_plan(select_plan)?;
     emit_program(database_header, optimized_plan, connection)
 }