From 9f726dbe62298f35bb0b96f30a2faf3559d8b050 Mon Sep 17 00:00:00 2001
From: pedrocarlo <pedro.muniz.carlo@gmail.com>
Date: Sat, 10 May 2025 22:36:43 -0300
Subject: [PATCH] simplify simple count detection

---
 core/translate/emitter.rs |  2 +-
 core/translate/plan.rs    | 44 +++++++++++++++++++++++++++++++++++++--
 core/translate/select.rs  | 44 ---------------------------------------
 3 files changed, 43 insertions(+), 47 deletions(-)

diff --git a/core/translate/emitter.rs b/core/translate/emitter.rs
index 65129a9f0..6557ca169 100644
--- a/core/translate/emitter.rs
+++ b/core/translate/emitter.rs
@@ -289,7 +289,7 @@ pub fn emit_query<'a>(
         OperationMode::SELECT,
     )?;
 
-    if plan.is_simple_count {
+    if plan.is_simple_count() {
         emit_simple_count(program, t_ctx, plan)?;
         return Ok(t_ctx.reg_result_cols_start.unwrap());
     }
diff --git a/core/translate/plan.rs b/core/translate/plan.rs
index 0ee865cb5..e34efb29b 100644
--- a/core/translate/plan.rs
+++ b/core/translate/plan.rs
@@ -308,8 +308,6 @@ 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 {
@@ -345,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)]
diff --git a/core/translate/select.rs b/core/translate/select.rs
index 90819169a..d6766a20f 100644
--- a/core/translate/select.rs
+++ b/core/translate/select.rs
@@ -106,7 +106,6 @@ 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();
@@ -388,8 +387,6 @@ 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))
         }
@@ -489,47 +486,6 @@ fn estimate_num_labels(select: &SelectPlan) -> usize {
     num_labels
 }
 
-/// 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(plan: &SelectPlan) -> bool {
-    if !plan.where_clause.is_empty()
-        || plan.aggregates.len() != 1
-        || matches!(plan.query_type, SelectQueryType::Subquery { .. })
-        || plan.table_references.len() != 1
-        || plan.result_columns.len() != 1
-        || plan.group_by.is_some()
-    // TODO: (pedrocarlo) maybe can optimize to use the count optmization with more columns
-    {
-        return false;
-    }
-    let table_ref = plan.table_references.first().unwrap();
-    if !matches!(table_ref.table, crate::schema::Table::BTree(..)) {
-        return false;
-    }
-    let agg = plan.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 = &plan.result_columns.get(0).unwrap().expr;
-    if *result_col_expr != count && *result_col_expr != count_star {
-        return false;
-    }
-    true
-}
-
 pub fn emit_simple_count<'a>(
     program: &mut ProgramBuilder,
     _t_ctx: &mut TranslateCtx<'a>,