Add countBy()

Author: antonmedv

bench_test.go

@@ -512,3 +512,17 @@ func Benchmark_groupBy(b *testing.B) {
 
 	require.Equal(b, 6, out.([]any)[0])
 }
+
+func Benchmark_countBy(b *testing.B) {
+	program, err := expr.Compile(`countBy(1..100, # % 7)[6]`)
+	require.NoError(b, err)
+
+	var out any
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		out, _ = vm.Run(program, nil)
+	}
+	b.StopTimer()
+
+	require.Equal(b, 14, out.(int))
+}

builtin/builtin.go

@@ -695,4 +695,9 @@ var Builtins = []*ast.Function{
 		Predicate: true,
 		Types:     types(new(func([]any, func(any) any) map[any][]any)),
 	},
+	{
+		Name:      "countBy",
+		Predicate: true,
+		Types:     types(new(func([]any, func(any) any) map[any]int)),
+	},
 }

checker/checker.go

@@ -736,6 +736,25 @@ func (v *checker) BuiltinNode(node *ast.BuiltinNode) (reflect.Type, info) {
 			return reflect.TypeOf(map[any][]any{}), info{}
 		}
 		return v.error(node.Arguments[1], "predicate should has one input and one output param")
+
+	case "countBy":
+		collection, _ := v.visit(node.Arguments[0])
+		if !isArray(collection) && !isAny(collection) {
+			return v.error(node.Arguments[0], "builtin %v takes only array (got %v)", node.Name, collection)
+		}
+
+		v.begin(collection)
+		closure, _ := v.visit(node.Arguments[1])
+		v.end()
+
+		if isFunc(closure) &&
+			closure.NumOut() == 1 &&
+			closure.NumIn() == 1 && isAny(closure.In(0)) {
+
+			return reflect.TypeOf(map[any]int{}), info{}
+		}
+		return v.error(node.Arguments[1], "predicate should has one input and one output param")
+
 	}
 
 	if id, ok := builtin.Index[node.Name]; ok {

compiler/compiler.go

@@ -824,6 +824,18 @@ func (c *compiler) BuiltinNode(node *ast.BuiltinNode) {
 		c.emit(OpGetGroupBy)
 		c.emit(OpEnd)
 		return
+
+	case "countBy":
+		c.compile(node.Arguments[0])
+		c.emit(OpBegin)
+		c.emitLoop(func() {
+			c.compile(node.Arguments[1])
+			c.emit(OpCountBy)
+		})
+		c.emit(OpGetCountBy)
+		c.emit(OpEnd)
+		return
+
 	}
 
 	if id, ok := builtin.Index[node.Name]; ok {

expr_test.go

@@ -1063,6 +1063,10 @@ func TestExpr(t *testing.T) {
 			`groupBy(ArrayOfFoo, .Value).foo`,
 			[]any{env.ArrayOfFoo[0]},
 		},
+		{
+			`countBy(1..9, # % 2)`,
+			map[any]int{0: 4, 1: 5},
+		},
 	}
 
 	for _, tt := range tests {

parser/parser.go

@@ -29,6 +29,7 @@ var predicates = map[string]struct {
 	"findLast":      {2},
 	"findLastIndex": {2},
 	"groupBy":       {2},
+	"countBy":       {2},
 }
 
 type parser struct {

vm/opcodes.go

@@ -73,9 +73,11 @@ const (
 	OpGetCount
 	OpGetLen
 	OpGetGroupBy
+	OpGetCountBy
 	OpPointer
 	OpThrow
 	OpGroupBy
+	OpCountBy
 	OpBegin
 	OpEnd // This opcode must be at the end of this list.
 )

vm/program.go

@@ -292,6 +292,9 @@ func (program *Program) Opcodes(w io.Writer) {
 		case OpGetGroupBy:
 			code("OpGetGroupBy")
 
+		case OpGetCountBy:
+			code("OpGetCountBy")
+
 		case OpPointer:
 			code("OpPointer")
 
@@ -304,6 +307,9 @@ func (program *Program) Opcodes(w io.Writer) {
 		case OpGroupBy:
 			code("OpGroupBy")
 
+		case OpCountBy:
+			code("OpCountBy")
+
 		case OpEnd:
 			code("OpEnd")
 

vm/vm.go

@@ -44,6 +44,7 @@ type Scope struct {
 	Len     int
 	Count   int
 	GroupBy map[any][]any
+	CountBy map[any]int
 }
 
 func Debug() *VM {
@@ -465,6 +466,9 @@ func (vm *VM) Run(program *Program, env any) (_ any, err error) {
 		case OpGetGroupBy:
 			vm.push(vm.Scope().GroupBy)
 
+		case OpGetCountBy:
+			vm.push(vm.Scope().CountBy)
+
 		case OpPointer:
 			scope := vm.Scope()
 			vm.push(scope.Array.Index(scope.Index).Interface())
@@ -481,6 +485,13 @@ func (vm *VM) Run(program *Program, env any) (_ any, err error) {
 			key := vm.pop()
 			scope.GroupBy[key] = append(scope.GroupBy[key], it)
 
+		case OpCountBy:
+			scope := vm.Scope()
+			if scope.CountBy == nil {
+				scope.CountBy = make(map[any]int)
+			}
+			scope.CountBy[vm.pop()]++
+
 		case OpBegin:
 			a := vm.pop()
 			array := reflect.ValueOf(a)