Add reduce()

Author: antonmedv

bench_test.go

@@ -526,3 +526,17 @@ func Benchmark_countBy(b *testing.B) {
 
 	require.Equal(b, 14, out.(int))
 }
+
+func Benchmark_reduce(b *testing.B) {
+	program, err := expr.Compile(`reduce(1..100, # + #acc)`)
+	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, 5050, out.(int))
+}

builtin/builtin.go

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

builtin/builtin_test.go

@@ -106,6 +106,10 @@ func TestBuiltin(t *testing.T) {
 		{`groupBy(1..3, # > 1 ? nil : "")[nil]`, []any{2, 3}},
 		{`groupBy(ArrayOfFoo, .Value).a`, []any{mock.Foo{Value: "a"}}},
 		{`countBy(1..9, # % 2)`, map[any]int{0: 4, 1: 5}},
+		{`reduce(1..9, # + #acc, 0)`, 45},
+		{`reduce(1..9, # + #acc)`, 45},
+		{`reduce([.5, 1.5, 2.5], # + #acc, 0)`, 4.5},
+		{`reduce([], 5, 0)`, 0},
 	}
 
 	for _, test := range tests {

checker/checker.go

@@ -755,6 +755,25 @@ func (v *checker) BuiltinNode(node *ast.BuiltinNode) (reflect.Type, info) {
 		}
 		return v.error(node.Arguments[1], "predicate should has one input and one output param")
 
+	case "reduce":
+		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, scopeVar{"index", integerType}, scopeVar{"acc", anyType})
+		closure, _ := v.visit(node.Arguments[1])
+		v.end()
+
+		if len(node.Arguments) == 3 {
+			_, _ = v.visit(node.Arguments[2])
+		}
+
+		if isFunc(closure) && closure.NumOut() == 1 {
+			return closure.Out(0), info{}
+		}
+		return v.error(node.Arguments[1], "predicate should has two input and one output param")
+
 	}
 
 	if id, ok := builtin.Index[node.Name]; ok {

compiler/compiler.go

@@ -836,6 +836,25 @@ func (c *compiler) BuiltinNode(node *ast.BuiltinNode) {
 		c.emit(OpEnd)
 		return
 
+	case "reduce":
+		c.compile(node.Arguments[0])
+		c.emit(OpBegin)
+		if len(node.Arguments) == 3 {
+			c.compile(node.Arguments[2])
+			c.emit(OpSetAcc)
+		} else {
+			c.emit(OpPointer)
+			c.emit(OpIncrementIndex)
+			c.emit(OpSetAcc)
+		}
+		c.emitLoop(func() {
+			c.compile(node.Arguments[1])
+			c.emit(OpSetAcc)
+		})
+		c.emit(OpGetAcc)
+		c.emit(OpEnd)
+		return
+
 	}
 
 	if id, ok := builtin.Index[node.Name]; ok {
@@ -903,6 +922,8 @@ func (c *compiler) PointerNode(node *ast.PointerNode) {
 	switch node.Name {
 	case "index":
 		c.emit(OpGetIndex)
+	case "acc":
+		c.emit(OpGetAcc)
 	case "":
 		c.emit(OpPointer)
 	default:

debug/debugger.go

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

parser/parser.go

@@ -30,6 +30,7 @@ var predicates = map[string]struct {
 	"findLastIndex": {2},
 	"groupBy":       {2},
 	"countBy":       {2},
+	"reduce":        {3},
 }
 
 type parser struct {
@@ -357,6 +358,9 @@ func (p *parser) parseCall(token Token) Node {
 
 		if b, ok := predicates[token.Value]; ok {
 			p.expect(Bracket, "(")
+
+			// TODO: Refactor parser to use builtin.Builtins instead of predicates map.
+
 			if b.arity == 1 {
 				arguments = make([]Node, 1)
 				arguments[0] = p.parseExpression(0)
@@ -366,6 +370,18 @@ func (p *parser) parseCall(token Token) Node {
 				p.expect(Operator, ",")
 				arguments[1] = p.parseClosure()
 			}
+
+			if token.Value == "reduce" {
+				arguments = make([]Node, 2)
+				arguments[0] = p.parseExpression(0)
+				p.expect(Operator, ",")
+				arguments[1] = p.parseClosure()
+				if p.current.Is(Operator, ",") {
+					p.next()
+					arguments = append(arguments, p.parseExpression(0))
+				}
+			}
+
 			p.expect(Bracket, ")")
 
 			node = &BuiltinNode{
@@ -596,9 +612,21 @@ func (p *parser) parsePipe(node Node) Node {
 
 	if b, ok := predicates[identifier.Value]; ok {
 		p.expect(Bracket, "(")
+
+		// TODO: Refactor parser to use builtin.Builtins instead of predicates map.
+
 		if b.arity == 2 {
 			arguments = append(arguments, p.parseClosure())
 		}
+
+		if identifier.Value == "reduce" {
+			arguments = append(arguments, p.parseClosure())
+			if p.current.Is(Operator, ",") {
+				p.next()
+				arguments = append(arguments, p.parseExpression(0))
+			}
+		}
+
 		p.expect(Bracket, ")")
 
 		node = &BuiltinNode{

vm/opcodes.go

@@ -74,10 +74,12 @@ const (
 	OpGetLen
 	OpGetGroupBy
 	OpGetCountBy
+	OpGetAcc
 	OpPointer
 	OpThrow
 	OpGroupBy
 	OpCountBy
+	OpSetAcc
 	OpBegin
 	OpEnd // This opcode must be at the end of this list.
 )

vm/program.go

@@ -295,21 +295,27 @@ func (program *Program) Opcodes(w io.Writer) {
 		case OpGetCountBy:
 			code("OpGetCountBy")
 
+		case OpGetAcc:
+			code("OpGetAcc")
+
 		case OpPointer:
 			code("OpPointer")
 
 		case OpThrow:
 			code("OpThrow")
 
-		case OpBegin:
-			code("OpBegin")
-
 		case OpGroupBy:
 			code("OpGroupBy")
 
 		case OpCountBy:
 			code("OpCountBy")
 
+		case OpSetAcc:
+			code("OpSetAcc")
+
+		case OpBegin:
+			code("OpBegin")
+
 		case OpEnd:
 			code("OpEnd")
 

vm/vm.go

@@ -45,6 +45,7 @@ type Scope struct {
 	Count   int
 	GroupBy map[any][]any
 	CountBy map[any]int
+	Acc     any
 }
 
 func Debug() *VM {
@@ -469,6 +470,12 @@ func (vm *VM) Run(program *Program, env any) (_ any, err error) {
 		case OpGetCountBy:
 			vm.push(vm.Scope().CountBy)
 
+		case OpGetAcc:
+			vm.push(vm.Scope().Acc)
+
+		case OpSetAcc:
+			vm.Scope().Acc = vm.pop()
+
 		case OpPointer:
 			scope := vm.Scope()
 			vm.push(scope.Array.Index(scope.Index).Interface())