Skip to content

Commit

Permalink
Add sum(map()) optimization
Browse files Browse the repository at this point in the history
  • Loading branch information
antonmedv committed Apr 13, 2024
1 parent 38f9496 commit edb1b5a
Show file tree
Hide file tree
Showing 4 changed files with 61 additions and 1 deletion.
1 change: 1 addition & 0 deletions optimizer/optimizer.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,5 +37,6 @@ func Optimize(node *Node, config *conf.Config) error {
Walk(node, &filterLast{})
Walk(node, &filterFirst{})
Walk(node, &predicateCombination{})
Walk(node, &sumMap{})
return nil
}
25 changes: 25 additions & 0 deletions optimizer/sum_map.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
package optimizer

import (
. "github.com/expr-lang/expr/ast"
)

type sumMap struct{}

func (*sumMap) Visit(node *Node) {
if sumBuiltin, ok := (*node).(*BuiltinNode); ok &&
sumBuiltin.Name == "sum" &&
len(sumBuiltin.Arguments) == 1 {
if mapBuiltin, ok := sumBuiltin.Arguments[0].(*BuiltinNode); ok &&
mapBuiltin.Name == "map" &&
len(mapBuiltin.Arguments) == 2 {
Patch(node, &BuiltinNode{
Name: "sum",
Arguments: []Node{
mapBuiltin.Arguments[0],
mapBuiltin.Arguments[1],
},
})
}
}
}
35 changes: 35 additions & 0 deletions optimizer/sum_map_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
package optimizer_test

import (
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

"github.com/expr-lang/expr/ast"
"github.com/expr-lang/expr/optimizer"
"github.com/expr-lang/expr/parser"
)

func TestOptimize_sum_map(t *testing.T) {
tree, err := parser.Parse(`sum(map(users, {.Age}))`)
require.NoError(t, err)

err = optimizer.Optimize(&tree.Node, nil)
require.NoError(t, err)

expected := &ast.BuiltinNode{
Name: "sum",
Arguments: []ast.Node{
&ast.IdentifierNode{Value: "users"},
&ast.ClosureNode{
Node: &ast.MemberNode{
Node: &ast.PointerNode{},
Property: &ast.StringNode{Value: "Age"},
},
},
},
}

assert.Equal(t, ast.Dump(expected), ast.Dump(tree.Node))
}
1 change: 0 additions & 1 deletion test/fuzz/fuzz_corpus.txt
Original file line number Diff line number Diff line change
Expand Up @@ -5608,7 +5608,6 @@ findLastIndex(list, f64 == f64)
findLastIndex(list, false == false)
findLastIndex(list, false and false)
findLastIndex(list, false and ok)
findLastIndex(list, false) != i
findLastIndex(list, foo != #)
findLastIndex(list, foo == foo)
findLastIndex(list, greet != half)
Expand Down

0 comments on commit edb1b5a

Please sign in to comment.