From 94358dc6659598731e042a091107785acda8e29f Mon Sep 17 00:00:00 2001
From: Vivek <vvk3785@gmail.com>
Date: Mon, 8 Jul 2024 13:55:06 +0530
Subject: [PATCH] Add aggregate fn: count

---
 core/translate.rs | 18 +++++++++++++++++-
 core/types.rs     |  3 +++
 core/vdbe.rs      | 14 ++++++++++++++
 testing/all.test  |  4 ++++
 4 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/core/translate.rs b/core/translate.rs
index 06b64d741..231b1630f 100644
--- a/core/translate.rs
+++ b/core/translate.rs
@@ -437,7 +437,23 @@ fn translate_aggregation(
             });
             target_register
         }
-        AggFunc::Count => todo!(),
+        AggFunc::Count => {
+            let expr_reg = if args.is_empty() {
+                program.alloc_register()
+            } else {
+                let expr = &args[0];
+                let expr_reg = program.alloc_register();
+                let _ = translate_expr(program, cursor_id, table, expr, expr_reg);
+                expr_reg
+            };
+            program.emit_insn(Insn::AggStep {
+                acc_reg: target_register,
+                col: expr_reg,
+                func: AggFunc::Count,
+            });
+            target_register
+        }
+        
         AggFunc::GroupConcat => todo!(),
         AggFunc::Max => todo!(),
         AggFunc::Min => todo!(),
diff --git a/core/types.rs b/core/types.rs
index 5d060ad57..fb35201c6 100644
--- a/core/types.rs
+++ b/core/types.rs
@@ -45,6 +45,7 @@ impl Display for OwnedValue {
             OwnedValue::Agg(a) => match a.as_ref() {
                 AggContext::Avg(acc, _count) => write!(f, "{}", acc),
                 AggContext::Sum(acc) => write!(f, "{}", acc),
+                AggContext::Count(count) => write!(f, "{}", count),
             },
         }
     }
@@ -54,6 +55,7 @@ impl Display for OwnedValue {
 pub enum AggContext {
     Avg(OwnedValue, OwnedValue), // acc and count
     Sum(OwnedValue),
+    Count(OwnedValue),
 }
 
 impl std::ops::Add<OwnedValue> for OwnedValue {
@@ -158,6 +160,7 @@ pub fn to_value(value: &OwnedValue) -> Value<'_> {
         OwnedValue::Agg(a) => match a.as_ref() {
             AggContext::Avg(acc, _count) => to_value(acc), // we assume aggfinal was called
             AggContext::Sum(acc) => to_value(acc),
+            AggContext::Count(count) => to_value(count),
         },
     }
 }
diff --git a/core/vdbe.rs b/core/vdbe.rs
index 76ac2aa78..cc8e76317 100644
--- a/core/vdbe.rs
+++ b/core/vdbe.rs
@@ -369,6 +369,9 @@ impl Program {
                             AggFunc::Sum => {
                                 OwnedValue::Agg(Box::new(AggContext::Sum(OwnedValue::Float(0.0))))
                             }
+                            AggFunc::Count => {
+                                OwnedValue::Agg(Box::new(AggContext::Count(OwnedValue::Integer(0))))
+                            }
                             _ => {
                                 todo!();
                             }
@@ -398,6 +401,16 @@ impl Program {
                             };
                             *acc += col;
                         }
+                        AggFunc::Count => {
+                            let OwnedValue::Agg(agg) = state.registers[*acc_reg].borrow_mut()
+                            else {
+                                unreachable!();
+                            };
+                            let AggContext::Count(count) = agg.borrow_mut() else {
+                                unreachable!();
+                            };
+                            *count += 1;
+                        }
                         _ => {
                             todo!();
                         }
@@ -417,6 +430,7 @@ impl Program {
                             *acc /= count.clone();
                         }
                         AggFunc::Sum => {}
+                        AggFunc::Count => {}
                         _ => {
                             todo!();
                         }
diff --git a/testing/all.test b/testing/all.test
index 7714788a8..f8cbaf647 100755
--- a/testing/all.test
+++ b/testing/all.test
@@ -43,6 +43,10 @@ do_execsql_test select-limit {
   SELECT id FROM users LIMIT 1;
 } {1}
 
+do_execsql_test select-count {
+  SELECT count(id) FROM users;
+} {10000}
+
 do_execsql_test pragma-cache-size {
   PRAGMA cache_size
 } {-2000}