Split Plan into Select and Delete

core/lib.rs

@@ -39,7 +39,7 @@ 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;
+use crate::translate::optimizer::optimize_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_select_plan(plan)?;
+                            let plan = optimize_plan(plan)?;
                             println!("{}", plan);
                         }
                         _ => todo!(),

core/translate/delete.rs

@@ -1,5 +1,5 @@
-use crate::translate::emitter::emit_program_for_delete;
-use crate::translate::optimizer::optimize_delete_plan;
+use crate::translate::emitter::emit_program;
+use crate::translate::optimizer::optimize_plan;
 use crate::translate::planner::prepare_delete_plan;
 use crate::{schema::Schema, storage::sqlite3_ondisk::DatabaseHeader, vdbe::Program};
 use crate::{Connection, Result};
@@ -16,6 +16,6 @@ pub fn translate_delete(
     connection: Weak<Connection>,
 ) -> Result<Program> {
     let delete_plan = prepare_delete_plan(schema, tbl_name, where_clause)?;
-    let optimized_plan = optimize_delete_plan(delete_plan)?;
-    emit_program_for_delete(database_header, optimized_plan, connection)
+    let optimized_plan = optimize_plan(delete_plan)?;
+    emit_program(database_header, optimized_plan, connection)
 }

core/translate/emitter.rs

@@ -9,7 +9,7 @@ use sqlite3_parser::ast::{self};
 
 use crate::schema::{Column, PseudoTable, Table};
 use crate::storage::sqlite3_ondisk::DatabaseHeader;
-use crate::translate::plan::{IterationDirection, Search};
+use crate::translate::plan::{DeletePlan, IterationDirection, Plan, Search};
 use crate::types::{OwnedRecord, OwnedValue};
 use crate::util::exprs_are_equivalent;
 use crate::vdbe::builder::ProgramBuilder;
@@ -20,7 +20,7 @@ use super::expr::{
     translate_aggregation, translate_aggregation_groupby, translate_condition_expr, translate_expr,
     ConditionMetadata,
 };
-use super::plan::{Aggregate, BTreeTableReference, Direction, GroupBy, Plan};
+use super::plan::{Aggregate, BTreeTableReference, Direction, GroupBy, SelectPlan};
 use super::plan::{ResultSetColumn, SourceOperator};
 
 // Metadata for handling LEFT JOIN operations
@@ -175,6 +175,17 @@ pub fn emit_program(
     database_header: Rc<RefCell<DatabaseHeader>>,
     mut plan: Plan,
     connection: Weak<Connection>,
+) -> Result<Program> {
+    match plan {
+        Plan::Select(plan) => emit_program_for_select(database_header, plan, connection),
+        Plan::Delete(plan) => emit_program_for_delete(database_header, plan, connection),
+    }
+}
+
+fn emit_program_for_select(
+    database_header: Rc<RefCell<DatabaseHeader>>,
+    mut plan: SelectPlan,
+    connection: Weak<Connection>,
 ) -> Result<Program> {
     let (mut program, mut metadata, init_label, start_offset) = prologue()?;
 
@@ -286,9 +297,9 @@ pub fn emit_program(
     Ok(program.build(database_header, connection))
 }
 
-pub fn emit_program_for_delete(
+fn emit_program_for_delete(
     database_header: Rc<RefCell<DatabaseHeader>>,
-    mut plan: Plan,
+    mut plan: DeletePlan,
     connection: Weak<Connection>,
 ) -> Result<Program> {
     let (mut program, mut metadata, init_label, start_offset) = prologue()?;
@@ -925,7 +936,7 @@ pub enum InnerLoopEmitTarget<'a> {
 /// At this point the cursors for all tables have been opened and rewound.
 fn inner_loop_emit(
     program: &mut ProgramBuilder,
-    plan: &mut Plan,
+    plan: &mut SelectPlan,
     metadata: &mut Metadata,
 ) -> Result<()> {
     // if we have a group by, we emit a record into the group by sorter.

core/translate/optimizer.rs

@@ -6,15 +6,22 @@ use crate::{schema::Index, Result};
 
 use super::plan::{
     get_table_ref_bitmask_for_ast_expr, get_table_ref_bitmask_for_operator, BTreeTableReference,
-    Direction, IterationDirection, Plan, Search, SourceOperator,
+    DeletePlan, Direction, IterationDirection, Plan, Search, SelectPlan, SourceOperator,
 };
 
+pub fn optimize_plan(mut plan: Plan) -> Result<Plan> {
+    match plan {
+        Plan::Select(plan) => optimize_select_plan(plan).map(Plan::Select),
+        Plan::Delete(plan) => optimize_delete_plan(plan).map(Plan::Delete),
+    }
+}
+
 /**
  * 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_select_plan(mut plan: Plan) -> Result<Plan> {
+fn optimize_select_plan(mut plan: SelectPlan) -> Result<SelectPlan> {
     eliminate_between(&mut plan.source, &mut plan.where_clause)?;
     if let ConstantConditionEliminationResult::ImpossibleCondition =
         eliminate_constants(&mut plan.source, &mut plan.where_clause)?
@@ -45,7 +52,7 @@ pub fn optimize_select_plan(mut plan: Plan) -> Result<Plan> {
     Ok(plan)
 }
 
-pub fn optimize_delete_plan(mut plan: Plan) -> Result<Plan> {
+fn optimize_delete_plan(mut plan: DeletePlan) -> Result<DeletePlan> {
     eliminate_between(&mut plan.source, &mut plan.where_clause)?;
     if let ConstantConditionEliminationResult::ImpossibleCondition =
         eliminate_constants(&mut plan.source, &mut plan.where_clause)?

core/translate/plan.rs

@@ -1,11 +1,12 @@
 use core::fmt;
+use sqlite3_parser::ast;
+use std::ptr::write;
 use std::{
     fmt::{Display, Formatter},
     rc::Rc,
 };
 
-use sqlite3_parser::ast;
-
+use crate::translate::plan::Plan::{Delete, Select};
 use crate::{
     function::AggFunc,
     schema::{BTreeTable, Column, Index},
@@ -27,7 +28,13 @@ pub struct GroupBy {
 }
 
 #[derive(Debug)]
-pub struct Plan {
+pub enum Plan {
+    Select(SelectPlan),
+    Delete(DeletePlan),
+}
+
+#[derive(Debug)]
+pub struct SelectPlan {
     /// A tree of sources (tables).
     pub source: SourceOperator,
     /// the columns inside SELECT ... FROM
@@ -50,9 +57,32 @@ pub struct Plan {
     pub contains_constant_false_condition: bool,
 }
 
+#[derive(Debug)]
+pub struct DeletePlan {
+    /// A tree of sources (tables).
+    pub source: SourceOperator,
+    /// the columns inside SELECT ... FROM
+    pub result_columns: Vec<ResultSetColumn>,
+    /// where clause split into a vec at 'AND' boundaries.
+    pub where_clause: Option<Vec<ast::Expr>>,
+    /// order by clause
+    pub order_by: Option<Vec<(ast::Expr, Direction)>>,
+    /// limit clause
+    pub limit: Option<usize>,
+    /// all the tables referenced in the query
+    pub referenced_tables: Vec<BTreeTableReference>,
+    /// all the indexes available
+    pub available_indexes: Vec<Rc<Index>>,
+    /// query contains a constant condition that is always false
+    pub contains_constant_false_condition: bool,
+}
+
 impl Display for Plan {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.source)
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Select(select_plan) => write!(f, "{}", select_plan.source),
+            Delete(delete_plan) => write!(f, "{}", delete_plan.source),
+        }
     }
 }
 

core/translate/planner.rs

@@ -1,5 +1,6 @@
 use super::plan::{
-    Aggregate, BTreeTableReference, Direction, GroupBy, Plan, ResultSetColumn, SourceOperator,
+    Aggregate, BTreeTableReference, DeletePlan, Direction, GroupBy, Plan, ResultSetColumn,
+    SelectPlan, SourceOperator,
 };
 use crate::{function::Func, schema::Schema, util::normalize_ident, Result};
 use sqlite3_parser::ast::{self, Expr, FromClause, JoinType, QualifiedName, ResultColumn};
@@ -269,7 +270,7 @@ pub fn prepare_select_plan<'a>(schema: &Schema, select: ast::Select) -> Result<P
             // Parse the FROM clause
             let (source, referenced_tables) = parse_from(schema, from, &mut operator_id_counter)?;
 
-            let mut plan = Plan {
+            let mut plan = SelectPlan {
                 source,
                 result_columns: vec![],
                 where_clause: None,
@@ -478,7 +479,7 @@ pub fn prepare_select_plan<'a>(schema: &Schema, select: ast::Select) -> Result<P
             }
 
             // Return the unoptimized query plan
-            Ok(plan)
+            Ok(Plan::Select(plan))
         }
         _ => todo!(),
     }
@@ -504,7 +505,7 @@ pub fn prepare_delete_plan(
     // Parse and resolve the where_clause
     let resolved_where_clauses = parse_where(where_clause, &[table_ref.clone()])?;
 
-    let plan = Plan {
+    let plan = DeletePlan {
         source: SourceOperator::Scan {
             id: 0,
             table_reference: table_ref.clone(),
@@ -513,16 +514,14 @@ pub fn prepare_delete_plan(
         },
         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)
+    Ok(Plan::Delete(plan))
 }
 
 #[allow(clippy::type_complexity)]

core/translate/select.rs

@@ -4,7 +4,7 @@ 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::translate::optimizer::optimize_plan;
 use crate::Connection;
 use crate::{schema::Schema, vdbe::Program, Result};
 use sqlite3_parser::ast;
@@ -16,6 +16,6 @@ pub fn translate_select(
     connection: Weak<Connection>,
 ) -> Result<Program> {
     let select_plan = prepare_select_plan(schema, select)?;
-    let optimized_plan = optimize_select_plan(select_plan)?;
+    let optimized_plan = optimize_plan(select_plan)?;
     emit_program(database_header, optimized_plan, connection)
 }