From 2cf27a1666695bdba90f7ff9a4ab45ddd73a5184 Mon Sep 17 00:00:00 2001
From: Vivek Patel <ervicpatelwork@gmail.com>
Date: Wed, 8 Sep 2021 08:33:41 +0530
Subject: [PATCH] Updates expr function count

- It returns count based on params
- count(foo) counts all elements
- count(foo, 'bar', 'baz') counts number of bars and bazs in foo
---
 pkg/expr/func_arr.go      | 26 ++++++++++++++++++++------
 pkg/expr/func_arr_test.go | 30 ++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/pkg/expr/func_arr.go b/pkg/expr/func_arr.go
index 8f9d5482b..3c275b29c 100644
--- a/pkg/expr/func_arr.go
+++ b/pkg/expr/func_arr.go
@@ -2,9 +2,9 @@ package expr
 
 import (
 	"fmt"
-	"reflect"
-
 	"github.com/PaesslerAG/gval"
+	"reflect"
+	"strings"
 )
 
 func ArrayFunctions() []gval.Language {
@@ -94,15 +94,29 @@ func shift(arr interface{}) (out interface{}, err error) {
 	return c.Index(0).Interface(), nil
 }
 
-// count gets the count of occurrences in the slice
+// count gets the count of occurrences in the first params(can be string, slice)
+// If given only one parameter then it returns length as count
 func count(arr interface{}, v ...interface{}) (count int, err error) {
+	var (
+		occ int
+		c   = reflect.ValueOf(arr)
+	)
+
+	if len(v) == 0 {
+		return c.Len(), nil
+	}
+
+	if c.Kind() == reflect.String {
+		for _, ww := range v {
+			count += strings.Count(c.String(), reflect.ValueOf(ww).String())
+		}
+		return
+	}
+
 	if arr, err = toSlice(arr); err != nil {
 		return
 	}
 
-	var (
-		occ int
-	)
 	for _, vv := range v {
 		if occ, err = find(arr, vv); err != nil {
 			return 0, err
diff --git a/pkg/expr/func_arr_test.go b/pkg/expr/func_arr_test.go
index 83a518a36..07906df20 100644
--- a/pkg/expr/func_arr_test.go
+++ b/pkg/expr/func_arr_test.go
@@ -264,6 +264,36 @@ func Test_count(t *testing.T) {
 			val:    []interface{}{0.1, 1.1},
 			expect: 0,
 		},
+		{
+			arr:    "foo",
+			val:    nil,
+			expect: 3,
+		},
+		{
+			arr:    "foo",
+			val:    []interface{}{},
+			expect: 3,
+		},
+		{
+			arr:    "foo bar",
+			val:    []interface{}{},
+			expect: 7,
+		},
+		{
+			arr:    []bool{true, true},
+			val:    []interface{}{},
+			expect: 2,
+		},
+		{
+			arr:    "foo",
+			val:    []interface{}{"bar", "baz"},
+			expect: 0,
+		},
+		{
+			arr:    "foo",
+			val:    []interface{}{"o", 12},
+			expect: 2,
+		},
 	}
 
 	for p, tc := range tcc {