fix is_simple_count detection

core/translate/emitter.rs

@@ -23,7 +23,7 @@ use super::group_by::{
 use super::main_loop::{close_loop, emit_loop, init_loop, open_loop, LeftJoinMetadata, LoopLabels};
 use super::order_by::{emit_order_by, init_order_by, SortMetadata};
 use super::plan::{JoinOrderMember, Operation, SelectPlan, TableReference, UpdatePlan};
-use super::select::{emit_simple_count, is_simple_count};
+use super::select::emit_simple_count;
 use super::schema::ParseSchema;
 use super::subquery::emit_subqueries;
 
@@ -233,8 +233,6 @@ pub fn emit_query<'a>(
     plan: &'a mut SelectPlan,
     t_ctx: &'a mut TranslateCtx<'a>,
 ) -> Result<usize> {
-    let is_simple_count = is_simple_count(plan);
-
     // Emit subqueries first so the results can be read in the main query loop.
     emit_subqueries(program, t_ctx, &mut plan.table_references)?;
 
@@ -312,7 +310,7 @@ pub fn emit_query<'a>(
         program.preassign_label_to_next_insn(jump_target_when_true);
     }
 
-    if !is_simple_count {
+    if !plan.is_simple_count {
         // Set up main query execution loop
         open_loop(
             program,
@@ -347,7 +345,7 @@ pub fn emit_query<'a>(
         group_by_emit_row_phase(program, t_ctx, plan)?;
     } else if !plan.aggregates.is_empty() {
         // Handle aggregation without GROUP BY
-        if is_simple_count {
+        if plan.is_simple_count {
             emit_simple_count(program, t_ctx, plan)?;
         } else {
             emit_ungrouped_aggregation(program, t_ctx, plan)?;

core/translate/plan.rs

@@ -308,6 +308,8 @@ pub struct SelectPlan {
     pub contains_constant_false_condition: bool,
     /// query type (top level or subquery)
     pub query_type: SelectQueryType,
+    /// if the query is of the format `SELECT count(*) FROM <tbl>`. This is set after the SelectPlan is create
+    pub is_simple_count: bool,
 }
 
 impl SelectPlan {

core/translate/select.rs

@@ -106,6 +106,7 @@ pub fn prepare_select_plan<'a>(
                 offset: None,
                 contains_constant_false_condition: false,
                 query_type: SelectQueryType::TopLevel,
+                is_simple_count: false,
             };
 
             let mut aggregate_expressions = Vec::new();
@@ -387,6 +388,8 @@ pub fn prepare_select_plan<'a>(
             (plan.limit, plan.offset) =
                 select.limit.map_or(Ok((None, None)), |l| parse_limit(&l))?;
 
+            plan.is_simple_count = is_simple_count(&plan);
+
             // Return the unoptimized query plan
             Ok(Plan::Select(plan))
         }
@@ -509,6 +512,22 @@ pub fn is_simple_count(plan: &SelectPlan) -> bool {
     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 = &plan.result_columns.get(0).unwrap().expr;
+    if *result_col_expr != count || *result_col_expr != count_star {
+        return false;
+    }
     true
 }