From c4e3cce8a2cc3dcd8327759b405c610f8a3c11df Mon Sep 17 00:00:00 2001
From: Pekka Enberg <penberg@iki.fi>
Date: Wed, 3 Jul 2024 21:54:25 +0300
Subject: [PATCH] Introduce Select struct

This introduces an intermediate `Select` struct, which hopefully makes
the codegen a bit simpler by transforming the complext AST to something
more straight-forward.
---
 core/translate.rs | 106 +++++++++++++++++++++++++++++-----------------
 1 file changed, 68 insertions(+), 38 deletions(-)

diff --git a/core/translate.rs b/core/translate.rs
index 07c8ff0b9..398dfeba5 100644
--- a/core/translate.rs
+++ b/core/translate.rs
@@ -2,13 +2,20 @@ use std::cell::RefCell;
 use std::rc::Rc;
 
 use crate::pager::Pager;
-use crate::schema::Schema;
+use crate::schema::{Schema, Table};
 use crate::sqlite3_ondisk::{DatabaseHeader, MIN_PAGE_CACHE_SIZE};
 use crate::util::normalize_ident;
 use crate::vdbe::{AggFunc, Insn, Program, ProgramBuilder};
 use anyhow::Result;
 use sqlite3_parser::ast;
 
+struct Select<'a> {
+    columns: Vec<ast::ResultColumn>,
+    column_info: Vec<ColumnInfo>,
+    from: Option<&'a Table>,
+    limit: Option<ast::Limit>,
+}
+
 enum AggregationFunc {
     Avg,
     Count,
@@ -48,14 +55,61 @@ pub fn translate(
     pager: Rc<Pager>,
 ) -> Result<Program> {
     match stmt {
-        ast::Stmt::Select(select) => translate_select(schema, select),
+        ast::Stmt::Select(select) => {
+            let select = build_select(schema, select)?;
+            translate_select(schema, select)
+        }
         ast::Stmt::Pragma(name, body) => translate_pragma(&name, body, database_header, pager),
         _ => todo!(),
     }
 }
 
+fn build_select(schema: &Schema, select: ast::Select) -> Result<Select> {
+    match select.body.select {
+        ast::OneSelect::Select {
+            columns,
+            from: Some(from),
+            ..
+        } => {
+            let table_name = match from.select {
+                Some(select_table) => match *select_table {
+                    ast::SelectTable::Table(name, ..) => name.name,
+                    _ => todo!(),
+                },
+                None => todo!(),
+            };
+            let table_name = table_name.0;
+            let table = match schema.get_table(&table_name) {
+                Some(table) => table,
+                None => anyhow::bail!("Parse error: no such table: {}", table_name),
+            };
+            let column_info = analyze_columns(&columns, Some(&table));
+            Ok(Select {
+                columns,
+                column_info,
+                from: Some(&table),
+                limit: select.limit.clone(),
+            })
+        }
+        ast::OneSelect::Select {
+            columns,
+            from: None,
+            ..
+        } => {
+            let column_info = analyze_columns(&columns, None);
+            Ok(Select {
+                columns,
+                column_info,
+                from: None,
+                limit: select.limit.clone(),
+            })
+        }
+        _ => todo!(),
+    }
+}
+
 /// Generate code for a SELECT statement.
-fn translate_select(schema: &Schema, select: ast::Select) -> Result<Program> {
+fn translate_select(schema: &Schema, select: Select) -> Result<Program> {
     let mut program = ProgramBuilder::new();
     let init_offset = program.emit_placeholder();
     let start_offset = program.offset();
@@ -72,25 +126,9 @@ fn translate_select(schema: &Schema, select: ast::Select) -> Result<Program> {
     } else {
         None
     };
-    let limit_decr_insn = match select.body.select {
-        ast::OneSelect::Select {
-            columns,
-            from: Some(from),
-            ..
-        } => {
-            let cursor_id = program.alloc_cursor_id();
-            let table_name = match from.select {
-                Some(select_table) => match *select_table {
-                    sqlite3_parser::ast::SelectTable::Table(name, ..) => name.name,
-                    _ => todo!(),
-                },
-                None => todo!(),
-            };
-            let table_name = table_name.0;
-            let table = match schema.get_table(&table_name) {
-                Some(table) => table,
-                None => anyhow::bail!("Parse error: no such table: {}", table_name),
-            };
+    let cursor_id = program.alloc_cursor_id();
+    let limit_decr_insn = match select.from {
+        Some(table) => {
             let root_page = table.root_page;
             program.emit_insn(Insn::OpenReadAsync {
                 cursor_id,
@@ -99,14 +137,13 @@ fn translate_select(schema: &Schema, select: ast::Select) -> Result<Program> {
             program.emit_insn(Insn::OpenReadAwait);
             program.emit_insn(Insn::RewindAsync { cursor_id });
             let rewind_await_offset = program.emit_placeholder();
-            let info_per_columns = analyze_columns(&columns, Some(table));
-            let exist_aggregation = info_per_columns.iter().any(|info| info.func.is_some());
+            let exist_aggregation = select.column_info.iter().any(|info| info.func.is_some());
             let (register_start, register_end) = translate_columns(
                 &mut program,
                 Some(cursor_id),
                 Some(table),
-                &columns,
-                &info_per_columns,
+                &select.columns,
+                &select.column_info,
                 exist_aggregation,
             );
             let limit_decr_insn = if exist_aggregation {
@@ -117,7 +154,7 @@ fn translate_select(schema: &Schema, select: ast::Select) -> Result<Program> {
                     pc_if_next: rewind_await_offset,
                 });
                 let mut target = register_start;
-                for info in &info_per_columns {
+                for info in &select.column_info {
                     if info.is_aggregation_function() {
                         let func = match info.func.as_ref().unwrap() {
                             AggregationFunc::Avg => AggFunc::Avg,
@@ -164,30 +201,23 @@ fn translate_select(schema: &Schema, select: ast::Select) -> Result<Program> {
             );
             limit_decr_insn
         }
-        ast::OneSelect::Select {
-            columns,
-            from: None,
-            ..
-        } => {
-            let info_per_columns = analyze_columns(&columns, None);
-            let exist_aggregation = info_per_columns.iter().any(|info| info.func.is_some());
+        None => {
+            let exist_aggregation = select.column_info.iter().any(|info| info.func.is_some());
             assert!(!exist_aggregation);
             let (register_start, register_end) = translate_columns(
                 &mut program,
                 None,
                 None,
-                &columns,
-                &info_per_columns,
+                &select.columns,
+                &select.column_info,
                 exist_aggregation,
             );
             program.emit_insn(Insn::ResultRow {
                 register_start,
                 register_end,
             });
-
             limit_reg.map(|_| program.emit_placeholder())
         }
-        _ => todo!(),
     };
     if let Some(limit_decr_insn) = limit_decr_insn {
         program.fixup_insn(