Remove TableReferenceType enum to clean up planner

tursodatabase · Feb 6, 2025 · b7bee6f · b7bee6f
1 parent 400b618
commit b7bee6f
Show file tree

Hide file tree

Showing 7 changed files with 74 additions and 100 deletions.
diff --git a/core/ext/mod.rs b/core/ext/mod.rs
@@ -136,6 +136,7 @@ impl Database {
             name: name.to_string(),
             implementation: vtab_module,
             columns,
+            args: None,
         };
         self.syms.borrow_mut().vtabs.insert(name.to_string(), vtab);
         ResultCode::OK

diff --git a/core/lib.rs b/core/lib.rs
@@ -515,6 +515,7 @@ impl<'a> Row<'a> {
 #[derive(Clone, Debug)]
 pub struct VirtualTable {
     name: String,
+    args: Option<Vec<ast::Expr>>,
     pub implementation: Rc<VTabModuleImpl>,
     columns: Vec<Column>,
 }

diff --git a/core/translate/delete.rs b/core/translate/delete.rs
@@ -7,7 +7,7 @@ use crate::vdbe::builder::{ProgramBuilder, ProgramBuilderOpts, QueryMode};
 use crate::{schema::Schema, Result, SymbolTable};
 use sqlite3_parser::ast::{Expr, Limit, QualifiedName};
 
-use super::plan::{TableReference, TableReferenceType};
+use super::plan::TableReference;
 
 pub fn translate_delete(
     query_mode: QueryMode,
@@ -48,7 +48,6 @@ pub fn prepare_delete_plan(
         identifier: table.name.clone(),
         op: Operation::Scan { iter_dir: None },
         join_info: None,
-        reference_type: TableReferenceType::BTreeTable,
     }];
 
     let mut where_predicates = vec![];

diff --git a/core/translate/expr.rs b/core/translate/expr.rs
@@ -3,7 +3,7 @@ use sqlite3_parser::ast::{self, UnaryOperator};
 #[cfg(feature = "json")]
 use crate::function::JsonFunc;
 use crate::function::{Func, FuncCtx, MathFuncArity, ScalarFunc};
-use crate::schema::Type;
+use crate::schema::{Table, Type};
 use crate::util::normalize_ident;
 use crate::vdbe::{
     builder::ProgramBuilder,
@@ -13,7 +13,7 @@ use crate::vdbe::{
 use crate::Result;
 
 use super::emitter::Resolver;
-use super::plan::{Operation, TableReference, TableReferenceType};
+use super::plan::{Operation, TableReference};
 
 #[derive(Debug, Clone, Copy)]
 pub struct ConditionMetadata {
@@ -1772,49 +1772,38 @@ pub fn translate_expr(
             match table_reference.op {
                 // If we are reading a column from a table, we find the cursor that corresponds to
                 // the table and read the column from the cursor.
-                Operation::Scan { .. } | Operation::Search(_) => {
-                    match &table_reference.reference_type {
-                        TableReferenceType::BTreeTable => {
-                            let cursor_id = program.resolve_cursor_id(&table_reference.identifier);
-                            if *is_rowid_alias {
-                                program.emit_insn(Insn::RowId {
-                                    cursor_id,
-                                    dest: target_register,
-                                });
-                            } else {
-                                program.emit_insn(Insn::Column {
-                                    cursor_id,
-                                    column: *column,
-                                    dest: target_register,
-                                });
-                            }
-                            let Some(column) = table_reference.table.get_column_at(*column) else {
-                                crate::bail_parse_error!("column index out of bounds");
-                            };
-                            maybe_apply_affinity(column.ty, target_register, program);
-                            Ok(target_register)
-                        }
-                        TableReferenceType::VirtualTable { .. } => {
-                            let cursor_id = program.resolve_cursor_id(&table_reference.identifier);
-                            program.emit_insn(Insn::VColumn {
+                Operation::Scan { .. } | Operation::Search(_) => match &table_reference.table {
+                    Table::BTree(_) => {
+                        let cursor_id = program.resolve_cursor_id(&table_reference.identifier);
+                        if *is_rowid_alias {
+                            program.emit_insn(Insn::RowId {
                                 cursor_id,
-                                column: *column,
                                 dest: target_register,
                             });
-                            Ok(target_register)
-                        }
-                        TableReferenceType::Subquery {
-                            result_columns_start_reg,
-                        } => {
-                            program.emit_insn(Insn::Copy {
-                                src_reg: result_columns_start_reg + *column,
-                                dst_reg: target_register,
-                                amount: 0,
+                        } else {
+                            program.emit_insn(Insn::Column {
+                                cursor_id,
+                                column: *column,
+                                dest: target_register,
                             });
-                            Ok(target_register)
                         }
+                        let Some(column) = table_reference.table.get_column_at(*column) else {
+                            crate::bail_parse_error!("column index out of bounds");
+                        };
+                        maybe_apply_affinity(column.ty, target_register, program);
+                        Ok(target_register)
                     }
-                }
+                    Table::Virtual(_) => {
+                        let cursor_id = program.resolve_cursor_id(&table_reference.identifier);
+                        program.emit_insn(Insn::VColumn {
+                            cursor_id,
+                            column: *column,
+                            dest: target_register,
+                        });
+                        Ok(target_register)
+                    }
+                    _ => unreachable!(),
+                },
                 // If we are reading a column from a subquery, we instead copy the column from the
                 // subquery's result registers.
                 Operation::Subquery {

diff --git a/core/translate/main_loop.rs b/core/translate/main_loop.rs
@@ -1,6 +1,7 @@
 use sqlite3_parser::ast;
 
 use crate::{
+    schema::Table,
     translate::result_row::emit_select_result,
     vdbe::{
         builder::{CursorType, ProgramBuilder},
@@ -17,7 +18,7 @@ use super::{
     order_by::{order_by_sorter_insert, sorter_insert},
     plan::{
         IterationDirection, Operation, Search, SelectPlan, SelectQueryType, TableReference,
-        TableReferenceType, WhereTerm,
+        WhereTerm,
     },
 };
 
@@ -78,37 +79,34 @@ pub fn init_loop(
         }
         match &table.op {
             Operation::Scan { .. } => {
-                let ref_type = &table.reference_type;
                 let cursor_id = program.alloc_cursor_id(
                     Some(table.identifier.clone()),
-                    match ref_type {
-                        TableReferenceType::BTreeTable => {
-                            CursorType::BTreeTable(table.btree().unwrap().clone())
-                        }
-                        TableReferenceType::VirtualTable { .. } => {
+                    match &table.table {
+                        Table::BTree(_) => CursorType::BTreeTable(table.btree().unwrap().clone()),
+                        Table::Virtual(_) => {
                             CursorType::VirtualTable(table.virtual_table().unwrap().clone())
                         }
                         other => panic!("Invalid table reference type in Scan: {:?}", other),
                     },
                 );
-                match (mode, ref_type) {
-                    (OperationMode::SELECT, TableReferenceType::BTreeTable) => {
+                match (mode, &table.table) {
+                    (OperationMode::SELECT, Table::BTree(_)) => {
                         let root_page = table.btree().unwrap().root_page;
                         program.emit_insn(Insn::OpenReadAsync {
                             cursor_id,
                             root_page,
                         });
                         program.emit_insn(Insn::OpenReadAwait {});
                     }
-                    (OperationMode::DELETE, TableReferenceType::BTreeTable) => {
+                    (OperationMode::DELETE, Table::BTree(_)) => {
                         let root_page = table.btree().unwrap().root_page;
                         program.emit_insn(Insn::OpenWriteAsync {
                             cursor_id,
                             root_page,
                         });
                         program.emit_insn(Insn::OpenWriteAwait {});
                     }
-                    (OperationMode::SELECT, TableReferenceType::VirtualTable { .. }) => {
+                    (OperationMode::SELECT, Table::Virtual(_)) => {
                         program.emit_insn(Insn::VOpenAsync { cursor_id });
                         program.emit_insn(Insn::VOpenAwait {});
                     }
@@ -258,10 +256,9 @@ pub fn open_loop(
                 }
             }
             Operation::Scan { iter_dir } => {
-                let ref_type = &table.reference_type;
                 let cursor_id = program.resolve_cursor_id(&table.identifier);
 
-                if !matches!(ref_type, TableReferenceType::VirtualTable { .. }) {
+                if !matches!(&table.table, Table::Virtual(_)) {
                     if iter_dir
                         .as_ref()
                         .is_some_and(|dir| *dir == IterationDirection::Backwards)
@@ -271,8 +268,8 @@ pub fn open_loop(
                         program.emit_insn(Insn::RewindAsync { cursor_id });
                     }
                 }
-                match ref_type {
-                    TableReferenceType::BTreeTable => program.emit_insn(
+                match &table.table {
+                    Table::BTree(_) => program.emit_insn(
                         if iter_dir
                             .as_ref()
                             .is_some_and(|dir| *dir == IterationDirection::Backwards)
@@ -288,13 +285,18 @@ pub fn open_loop(
                             }
                         },
                     ),
-                    TableReferenceType::VirtualTable { args, .. } => {
+                    Table::Virtual(ref table) => {
+                        let args = if let Some(args) = table.args.as_ref() {
+                            args
+                        } else {
+                            &vec![]
+                        };
                         let start_reg = program.alloc_registers(args.len());
                         let mut cur_reg = start_reg;
                         for arg in args {
                             let reg = cur_reg;
                             cur_reg += 1;
-                            translate_expr(program, Some(tables), arg, reg, &t_ctx.resolver)?;
+                            translate_expr(program, Some(tables), &arg, reg, &t_ctx.resolver)?;
                         }
                         program.emit_insn(Insn::VFilter {
                             cursor_id,
@@ -722,11 +724,10 @@ pub fn close_loop(
                 });
             }
             Operation::Scan { iter_dir, .. } => {
-                let ref_type = &table.reference_type;
                 program.resolve_label(loop_labels.next, program.offset());
                 let cursor_id = program.resolve_cursor_id(&table.identifier);
-                match ref_type {
-                    TableReferenceType::BTreeTable { .. } => {
+                match &table.table {
+                    Table::BTree(_) => {
                         if iter_dir
                             .as_ref()
                             .is_some_and(|dir| *dir == IterationDirection::Backwards)
@@ -750,7 +751,7 @@ pub fn close_loop(
                             });
                         }
                     }
-                    TableReferenceType::VirtualTable { .. } => {
+                    Table::Virtual(_) => {
                         program.emit_insn(Insn::VNext {
                             cursor_id,
                             pc_if_next: loop_labels.loop_start,

diff --git a/core/translate/plan.rs b/core/translate/plan.rs
@@ -198,7 +198,6 @@ pub struct TableReference {
     pub identifier: String,
     /// The join info for this table reference, if it is the right side of a join (which all except the first table reference have)
     pub join_info: Option<JoinInfo>,
-    pub reference_type: TableReferenceType,
 }
 
 #[derive(Clone, Debug)]
@@ -225,35 +224,17 @@ pub enum Operation {
     },
 }
 
-/// The type of the table reference, either BTreeTable or Subquery
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub enum TableReferenceType {
-    /// A BTreeTable is a table that is stored on disk in a B-tree index.
-    BTreeTable,
-    /// A subquery.
-    Subquery {
-        /// The index of the first register in the query plan that contains the result columns of the subquery.
-        result_columns_start_reg: usize,
-    },
-    /// A virtual table.
-    VirtualTable {
-        /// Arguments to pass e.g. generate_series(1, 10, 2)
-        args: Vec<ast::Expr>,
-    },
-}
-
 impl TableReference {
     /// Returns the btree table for this table reference, if it is a BTreeTable.
     pub fn btree(&self) -> Option<Rc<BTreeTable>> {
-        match &self.reference_type {
-            TableReferenceType::BTreeTable => self.table.btree(),
-            TableReferenceType::Subquery { .. } => None,
-            TableReferenceType::VirtualTable { .. } => None,
+        match &self.table {
+            Table::BTree(_) => self.table.btree(),
+            _ => None,
         }
     }
     pub fn virtual_table(&self) -> Option<Rc<VirtualTable>> {
-        match &self.reference_type {
-            TableReferenceType::VirtualTable { .. } => self.table.virtual_table(),
+        match &self.table {
+            Table::Virtual(_) => self.table.virtual_table(),
             _ => None,
         }
     }
@@ -280,9 +261,6 @@ impl TableReference {
                 result_columns_start_reg: 0, // Will be set in the bytecode emission phase
             },
             table,
-            reference_type: TableReferenceType::Subquery {
-                result_columns_start_reg: 0, // Will be set in the bytecode emission phase
-            },
             identifier: identifier.clone(),
             join_info,
         }

diff --git a/core/translate/planner.rs b/core/translate/planner.rs
@@ -1,7 +1,7 @@
 use super::{
     plan::{
         Aggregate, JoinInfo, Operation, Plan, ResultSetColumn, SelectQueryType, TableReference,
-        TableReferenceType, WhereTerm,
+        WhereTerm,
     },
     select::prepare_select_plan,
     SymbolTable,
@@ -11,7 +11,7 @@ use crate::{
     schema::{Schema, Table},
     util::{exprs_are_equivalent, normalize_ident},
     vdbe::BranchOffset,
-    Result,
+    Result, VirtualTable,
 };
 use sqlite3_parser::ast::{self, Expr, FromClause, JoinType, Limit, UnaryOperator};
 
@@ -301,7 +301,6 @@ fn parse_from_clause_table(
                 table: Table::BTree(table.clone()),
                 identifier: alias.unwrap_or(normalized_qualified_name),
                 join_info: None,
-                reference_type: TableReferenceType::BTreeTable,
             })
         }
         ast::SelectTable::Select(subselect, maybe_alias) => {
@@ -320,9 +319,9 @@ fn parse_from_clause_table(
                 .unwrap_or(format!("subquery_{}", cur_table_index));
             Ok(TableReference::new_subquery(identifier, subplan, None))
         }
-        ast::SelectTable::TableCall(qualified_name, mut maybe_args, maybe_alias) => {
-            let normalized_name = normalize_ident(qualified_name.name.0.as_str());
-            let Some(vtab) = syms.vtabs.get(&normalized_name) else {
+        ast::SelectTable::TableCall(qualified_name, maybe_args, maybe_alias) => {
+            let normalized_name = &normalize_ident(qualified_name.name.0.as_str());
+            let Some(vtab) = syms.vtabs.get(normalized_name) else {
                 crate::bail_parse_error!("Virtual table {} not found", normalized_name);
             };
             let alias = maybe_alias
@@ -331,16 +330,22 @@ fn parse_from_clause_table(
                     ast::As::As(id) => id.0.clone(),
                     ast::As::Elided(id) => id.0.clone(),
                 })
-                .unwrap_or(normalized_name);
+                .unwrap_or(normalized_name.to_string());
 
             Ok(TableReference {
                 op: Operation::Scan { iter_dir: None },
                 join_info: None,
-                table: Table::Virtual(vtab.clone().into()),
+                table: Table::Virtual(
+                    VirtualTable {
+                        name: normalized_name.clone(),
+                        args: maybe_args,
+                        implementation: vtab.implementation.clone(),
+                        columns: vtab.columns.clone(),
+                    }
+                    .into(),
+                )
+                .into(),
                 identifier: alias.clone(),
-                reference_type: TableReferenceType::VirtualTable {
-                    args: maybe_args.take().unwrap_or_default(),
-                },
             })
         }
         _ => todo!(),