From 725b7021c9f69b7aa6d37cd9502cab3d0ca2ab9d Mon Sep 17 00:00:00 2001
From: Anton Medvedev <anton@medv.io>
Date: Sat, 13 Apr 2024 23:38:35 +0200
Subject: [PATCH] Make count() predicate optional

---
 checker/checker.go          |  4 ++++
 compiler/compiler.go        |  6 +++++-
 docs/language-definition.md | 31 ++++++++++++++++++++++++++++---
 expr_test.go                |  4 ++++
 parser/parser.go            |  2 +-
 5 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/checker/checker.go b/checker/checker.go
index a6daa27b..ecf7a04d 100644
--- a/checker/checker.go
+++ b/checker/checker.go
@@ -653,6 +653,10 @@ func (v *checker) BuiltinNode(node *ast.BuiltinNode) (reflect.Type, info) {
 			return v.error(node.Arguments[0], "builtin %v takes only array (got %v)", node.Name, collection)
 		}
 
+		if len(node.Arguments) == 1 {
+			return integerType, info{}
+		}
+
 		v.begin(collection)
 		closure, _ := v.visit(node.Arguments[1])
 		v.end()
diff --git a/compiler/compiler.go b/compiler/compiler.go
index 07bc58b8..05a92f2f 100644
--- a/compiler/compiler.go
+++ b/compiler/compiler.go
@@ -800,7 +800,11 @@ func (c *compiler) BuiltinNode(node *ast.BuiltinNode) {
 		c.compile(node.Arguments[0])
 		c.emit(OpBegin)
 		c.emitLoop(func() {
-			c.compile(node.Arguments[1])
+			if len(node.Arguments) == 2 {
+				c.compile(node.Arguments[1])
+			} else {
+				c.emit(OpPointer)
+			}
 			c.emitCond(func() {
 				c.emit(OpIncrementCount)
 			})
diff --git a/docs/language-definition.md b/docs/language-definition.md
index a4e1a7cb..0bc2bd0d 100644
--- a/docs/language-definition.md
+++ b/docs/language-definition.md
@@ -628,14 +628,24 @@ Groups the elements of an array by the result of the [predicate](#predicate).
 groupBy(users, .Age)
 ```
 
-### count(array, predicate) {#count}
+### count(array[, predicate]) {#count}
 
 Returns the number of elements what satisfies the [predicate](#predicate).
 
+```expr
+count(users, .Age > 18)
+```
+
 Equivalent to:
 
 ```expr
-len(filter(array, predicate))
+len(filter(users, .Age > 18))
+```
+
+If the predicate is not given, returns the number of `true` elements in the array.
+
+```expr
+count([true, false, true]) == 2
 ```
 
 ### concat(array1, array2[, ...]) {#concat}
@@ -673,7 +683,7 @@ reduce(1..9, #acc + #)
 reduce(1..9, #acc + #, 0)
 ```
 
-### sum(array) {#sum}
+### sum(array[, predicate]) {#sum}
 
 Returns the sum of all numbers in the array.
 
@@ -681,6 +691,21 @@ Returns the sum of all numbers in the array.
 sum([1, 2, 3]) == 6
 ```
 
+If the optional `predicate` argument is given, it is a predicate that is applied on each element 
+of the array before summing.
+
+```expr
+sum(accounts, .Balance)
+```
+
+Equivalent to:
+
+```expr
+reduce(accounts, #acc + .Balance, 0)
+// or
+sum(map(accounts, .Balance))
+```
+
 ### mean(array) {#mean}
 
 Returns the average of all numbers in the array.
diff --git a/expr_test.go b/expr_test.go
index 38b97eaa..23f4c496 100644
--- a/expr_test.go
+++ b/expr_test.go
@@ -905,6 +905,10 @@ func TestExpr(t *testing.T) {
 			`count(1..30, {# % 3 == 0})`,
 			10,
 		},
+		{
+			`count([true, true, false])`,
+			2,
+		},
 		{
 			`"a" < "b"`,
 			true,
diff --git a/parser/parser.go b/parser/parser.go
index 641369b1..6d96561a 100644
--- a/parser/parser.go
+++ b/parser/parser.go
@@ -33,7 +33,7 @@ var predicates = map[string]struct {
 	"one":           {[]arg{expr, closure}},
 	"filter":        {[]arg{expr, closure}},
 	"map":           {[]arg{expr, closure}},
-	"count":         {[]arg{expr, closure}},
+	"count":         {[]arg{expr, closure | optional}},
 	"sum":           {[]arg{expr, closure | optional}},
 	"find":          {[]arg{expr, closure}},
 	"findIndex":     {[]arg{expr, closure}},