Add groupBy

Author: antonmedv

builtin/builtin.go

@@ -690,4 +690,9 @@ var Builtins = []*ast.Function{
 			return arrayType, nil
 		},
 	},
+	{
+		Name:      "groupBy",
+		Predicate: true,
+		Types:     types(new(func([]any, func(any) any) map[any][]any)),
+	},
 }

checker/checker.go

@@ -718,6 +718,24 @@ func (v *checker) BuiltinNode(node *ast.BuiltinNode) (reflect.Type, info) {
 			return integerType, info{}
 		}
 		return v.error(node.Arguments[1], "predicate should has one input and one output param")
+
+	case "groupBy":
+		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][]any{}), 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

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

debug/debugger.go

@@ -130,6 +130,9 @@ func StartDebugger(program *Program, env interface{}) {
 				keys = append(keys, pair{"Index", s.Index})
 				keys = append(keys, pair{"Len", s.Len})
 				keys = append(keys, pair{"Count", s.Count})
+				if s.GroupBy != nil {
+					keys = append(keys, pair{"GroupBy", s.GroupBy})
+				}
 				row := 0
 				for _, pair := range keys {
 					scope.SetCellSimple(row, 0, fmt.Sprintf("%v: ", pair.key))

expr_test.go

@@ -1040,6 +1040,29 @@ func TestExpr(t *testing.T) {
 			`4/2 == 2`,
 			true,
 		},
+		{
+			`groupBy(1..9, # % 2)`,
+			map[any][]any{
+				0: {2, 4, 6, 8},
+				1: {1, 3, 5, 7, 9},
+			},
+		},
+		{
+			`groupBy(1..9, # % 2)[0]`,
+			[]any{2, 4, 6, 8},
+		},
+		{
+			`groupBy(1..3, # > 1)[true]`,
+			[]any{2, 3},
+		},
+		{
+			`groupBy(1..3, # > 1 ? nil : "")[nil]`,
+			[]any{2, 3},
+		},
+		{
+			`groupBy(ArrayOfFoo, .Value).foo`,
+			[]any{env.ArrayOfFoo[0]},
+		},
 	}
 
 	for _, tt := range tests {

parser/parser.go

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

vm/opcodes.go

@@ -72,8 +72,10 @@ const (
 	OpSetIndex
 	OpGetCount
 	OpGetLen
+	OpGetGroupBy
 	OpPointer
 	OpThrow
+	OpGroupBy
 	OpBegin
 	OpEnd // This opcode must be at the end of this list.
 )

vm/program.go

@@ -289,6 +289,9 @@ func (program *Program) Opcodes(w io.Writer) {
 		case OpGetLen:
 			code("OpGetLen")
 
+		case OpGetGroupBy:
+			code("OpGetGroupBy")
+
 		case OpPointer:
 			code("OpPointer")
 
@@ -298,6 +301,9 @@ func (program *Program) Opcodes(w io.Writer) {
 		case OpBegin:
 			code("OpBegin")
 
+		case OpGroupBy:
+			code("OpGroupBy")
+
 		case OpEnd:
 			code("OpEnd")
 

vm/vm.go

@@ -39,10 +39,11 @@ type VM struct {
 }
 
 type Scope struct {
-	Array reflect.Value
-	Index int
-	Len   int
-	Count int
+	Array   reflect.Value
+	Index   int
+	Len     int
+	Count   int
+	GroupBy map[any][]any
 }
 
 func Debug() *VM {
@@ -461,13 +462,25 @@ func (vm *VM) Run(program *Program, env interface{}) (_ interface{}, err error)
 			scope := vm.Scope()
 			vm.push(scope.Len)
 
+		case OpGetGroupBy:
+			vm.push(vm.Scope().GroupBy)
+
 		case OpPointer:
 			scope := vm.Scope()
 			vm.push(scope.Array.Index(scope.Index).Interface())
 
 		case OpThrow:
 			panic(vm.pop().(error))
 
+		case OpGroupBy:
+			scope := vm.Scope()
+			if scope.GroupBy == nil {
+				scope.GroupBy = make(map[any][]any)
+			}
+			it := scope.Array.Index(scope.Index).Interface()
+			key := vm.pop()
+			scope.GroupBy[key] = append(scope.GroupBy[key], it)
+
 		case OpBegin:
 			a := vm.pop()
 			array := reflect.ValueOf(a)