Merge 'Count optimization' from Pedro Muniz

After reading #1123, I wanted to see what optimizations I could do.
Sqlite optimizes `count` aggregation for the following case: `SELECT
count() FROM <tbl>`. This is so widely used, that they made an
optimization just for it in the form of the `COUNT` opcode.
This PR thus implements this optimization by creating the `COUNT`
opcode, and checking in the select emitter if we the query is a Simple
Count Query. If it is, we just emit the Opcode instead of going through
a Rewind loop, saving on execution time.
The screenshots below show a huge decrease in execution time.
- **Main**
<img width="383" alt="image" src="https://github.com/user-
attachments/assets/99a9dec4-e7c5-41db-ba67-4eafa80dd2e6" />
- **Count Optimization**
<img width="435" alt="image" src="https://github.com/user-
attachments/assets/e93b3233-92e6-4736-aa60-b52b2477179f" />

Reviewed-by: Jussi Saurio <jussi.saurio@gmail.com>

Closes #1443

core/translate/emitter.rs

@@ -24,6 +24,7 @@ use super::main_loop::{close_loop, emit_loop, init_loop, open_loop, LeftJoinMeta
 use super::order_by::{emit_order_by, init_order_by, SortMetadata};
 use super::plan::{JoinOrderMember, Operation, SelectPlan, TableReference, UpdatePlan};
 use super::schema::ParseSchema;
+use super::select::emit_simple_count;
 use super::subquery::emit_subqueries;
 
 #[derive(Debug)]
@@ -288,6 +289,11 @@ pub fn emit_query<'a>(
         OperationMode::SELECT,
     )?;
 
+    if plan.is_simple_count() {
+        emit_simple_count(program, t_ctx, plan)?;
+        return Ok(t_ctx.reg_result_cols_start.unwrap());
+    }
+
     for where_term in plan
         .where_clause
         .iter()
@@ -323,6 +329,7 @@ pub fn emit_query<'a>(
 
     // Clean up and close the main execution loop
     close_loop(program, t_ctx, &plan.table_references, &plan.join_order)?;
+
     program.preassign_label_to_next_insn(after_main_loop_label);
 
     let mut order_by_necessary = plan.order_by.is_some() && !plan.contains_constant_false_condition;

core/translate/plan.rs

@@ -343,6 +343,48 @@ impl SelectPlan {
     pub fn group_by_sorter_column_count(&self) -> usize {
         self.agg_args_count() + self.group_by_col_count() + self.non_group_by_non_agg_column_count()
     }
+
+    /// Reference: https://github.com/sqlite/sqlite/blob/5db695197b74580c777b37ab1b787531f15f7f9f/src/select.c#L8613
+    ///
+    /// Checks to see if the query is of the format `SELECT count(*) FROM <tbl>`
+    pub fn is_simple_count(&self) -> bool {
+        if !self.where_clause.is_empty()
+            || self.aggregates.len() != 1
+            || matches!(self.query_type, SelectQueryType::Subquery { .. })
+            || self.table_references.len() != 1
+            || self.result_columns.len() != 1
+            || self.group_by.is_some()
+            || self.contains_constant_false_condition
+        // TODO: (pedrocarlo) maybe can optimize to use the count optmization with more columns
+        {
+            return false;
+        }
+        let table_ref = self.table_references.first().unwrap();
+        if !matches!(table_ref.table, crate::schema::Table::BTree(..)) {
+            return false;
+        }
+        let agg = self.aggregates.first().unwrap();
+        if !matches!(agg.func, AggFunc::Count0) {
+            return false;
+        }
+
+        let count = limbo_sqlite3_parser::ast::Expr::FunctionCall {
+            name: limbo_sqlite3_parser::ast::Id("count".to_string()),
+            distinctness: None,
+            args: None,
+            order_by: None,
+            filter_over: None,
+        };
+        let count_star = limbo_sqlite3_parser::ast::Expr::FunctionCallStar {
+            name: limbo_sqlite3_parser::ast::Id("count".to_string()),
+            filter_over: None,
+        };
+        let result_col_expr = &self.result_columns.get(0).unwrap().expr;
+        if *result_col_expr != count && *result_col_expr != count_star {
+            return false;
+        }
+        true
+    }
 }
 
 #[allow(dead_code)]

core/translate/select.rs

@@ -1,4 +1,4 @@
-use super::emitter::emit_program;
+use super::emitter::{emit_program, TranslateCtx};
 use super::plan::{select_star, JoinOrderMember, Operation, Search, SelectQueryType};
 use super::planner::Scope;
 use crate::function::{AggFunc, ExtFunc, Func};
@@ -10,6 +10,7 @@ use crate::translate::planner::{
 };
 use crate::util::normalize_ident;
 use crate::vdbe::builder::{ProgramBuilderOpts, QueryMode};
+use crate::vdbe::insn::Insn;
 use crate::SymbolTable;
 use crate::{schema::Schema, vdbe::builder::ProgramBuilder, Result};
 use limbo_sqlite3_parser::ast::{self, SortOrder};
@@ -484,3 +485,41 @@ fn estimate_num_labels(select: &SelectPlan) -> usize {
 
     num_labels
 }
+
+pub fn emit_simple_count<'a>(
+    program: &mut ProgramBuilder,
+    _t_ctx: &mut TranslateCtx<'a>,
+    plan: &'a SelectPlan,
+) -> Result<()> {
+    let cursors = plan
+        .table_references
+        .get(0)
+        .unwrap()
+        .resolve_cursors(program)?;
+
+    let cursor_id = {
+        match cursors {
+            (_, Some(cursor_id)) | (Some(cursor_id), None) => cursor_id,
+            _ => panic!("cursor for table should have been opened"),
+        }
+    };
+
+    // TODO: I think this allocation can be avoided if we are smart with the `TranslateCtx`
+    let target_reg = program.alloc_register();
+
+    program.emit_insn(Insn::Count {
+        cursor_id,
+        target_reg,
+        exact: true,
+    });
+
+    program.emit_insn(Insn::Close { cursor_id });
+    let output_reg = program.alloc_register();
+    program.emit_insn(Insn::Copy {
+        src_reg: target_reg,
+        dst_reg: output_reg,
+        amount: 0,
+    });
+    program.emit_result_row(output_reg, 1);
+    Ok(())
+}