Add filter()[0] optimization

Author: antonmedv

ast/node.go

@@ -133,6 +133,7 @@ type BuiltinNode struct {
 	base
 	Name      string
 	Arguments []Node
+	Throws    bool
 }
 
 type ClosureNode struct {

compiler/compiler.go

@@ -733,7 +733,12 @@ func (c *compiler) BuiltinNode(node *ast.BuiltinNode) {
 			c.patchJump(noop)
 			c.emit(OpPop)
 		})
-		c.emit(OpNil)
+		if node.Throws {
+			c.emit(OpPush, c.addConstant(fmt.Errorf("reflect: slice index out of range")))
+			c.emit(OpThrow)
+		} else {
+			c.emit(OpNil)
+		}
 		c.patchJump(loopBreak)
 		c.emit(OpEnd)
 		return
@@ -751,7 +756,7 @@ func (c *compiler) BuiltinNode(node *ast.BuiltinNode) {
 			c.patchJump(noop)
 			c.emit(OpPop)
 		})
-		c.emit(OpPushInt, -1)
+		c.emit(OpNil)
 		c.patchJump(loopBreak)
 		c.emit(OpEnd)
 		return

expr_test.go

@@ -972,6 +972,18 @@ func TestExpr(t *testing.T) {
 			`findIndex(ArrayOfFoo, .Value == "baz")`,
 			2,
 		},
+		{
+			`filter(ArrayOfFoo, .Value == "baz")[0]`,
+			env.ArrayOfFoo[2],
+		},
+		{
+			`first(filter(ArrayOfFoo, .Value == "baz"))`,
+			env.ArrayOfFoo[2],
+		},
+		{
+			`first(filter(ArrayOfFoo, false))`,
+			nil,
+		},
 	}
 
 	for _, tt := range tests {
@@ -1003,6 +1015,33 @@ func TestExpr(t *testing.T) {
 	}
 }
 
+func TestExpr_error(t *testing.T) {
+	env := mock.Env{}
+
+	tests := []struct {
+		code string
+		want string
+	}{
+		{
+			`filter(1..9, # > 9)[0]`,
+			`reflect: slice index out of range (1:20)
+ | filter(1..9, # > 9)[0]
+ | ...................^`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.code, func(t *testing.T) {
+			program, err := expr.Compile(tt.code, expr.Env(mock.Env{}))
+			require.NoError(t, err)
+
+			_, err = expr.Run(program, env)
+			require.Error(t, err)
+			assert.Equal(t, tt.want, err.Error())
+		})
+	}
+}
+
 func TestExpr_optional_chaining(t *testing.T) {
 	env := map[string]interface{}{}
 	program, err := expr.Compile("foo?.bar.baz", expr.Env(env), expr.AllowUndefinedVariables())

optimizer/filter_first.go

@@ -0,0 +1,36 @@
+package optimizer
+
+import (
+	. "github.com/antonmedv/expr/ast"
+)
+
+type filterFirst struct{}
+
+func (*filterFirst) Visit(node *Node) {
+	if member, ok := (*node).(*MemberNode); ok && member.Property != nil && !member.Optional {
+		if prop, ok := member.Property.(*IntegerNode); ok && prop.Value == 0 {
+			if filter, ok := member.Node.(*BuiltinNode); ok &&
+				filter.Name == "filter" &&
+				len(filter.Arguments) == 2 {
+				Patch(node, &BuiltinNode{
+					Name:      "find",
+					Arguments: filter.Arguments,
+					Throws:    true, // to match the behavior of filter()[0]
+				})
+			}
+		}
+	}
+	if first, ok := (*node).(*BuiltinNode); ok &&
+		first.Name == "first" &&
+		len(first.Arguments) == 1 {
+		if filter, ok := first.Arguments[0].(*BuiltinNode); ok &&
+			filter.Name == "filter" &&
+			len(filter.Arguments) == 2 {
+			Patch(node, &BuiltinNode{
+				Name:      "find",
+				Arguments: filter.Arguments,
+				Throws:    false, // as first() will return nil if not found
+			})
+		}
+	}
+}

optimizer/optimizer.go

@@ -34,5 +34,6 @@ func Optimize(node *Node, config *conf.Config) error {
 	Walk(node, &inRange{})
 	Walk(node, &constRange{})
 	Walk(node, &filterLen{})
+	Walk(node, &filterFirst{})
 	return nil
 }

optimizer/optimizer_test.go

@@ -165,3 +165,59 @@ func TestOptimize_filter_len(t *testing.T) {
 
 	assert.Equal(t, ast.Dump(expected), ast.Dump(tree.Node))
 }
+
+func TestOptimize_filter_0(t *testing.T) {
+	tree, err := parser.Parse(`filter(users, .Name == "Bob")[0]`)
+	require.NoError(t, err)
+
+	err = optimizer.Optimize(&tree.Node, nil)
+	require.NoError(t, err)
+
+	expected := &ast.BuiltinNode{
+		Name: "find",
+		Arguments: []ast.Node{
+			&ast.IdentifierNode{Value: "users"},
+			&ast.ClosureNode{
+				Node: &ast.BinaryNode{
+					Operator: "==",
+					Left: &ast.MemberNode{
+						Node:     &ast.PointerNode{},
+						Property: &ast.StringNode{Value: "Name"},
+					},
+					Right: &ast.StringNode{Value: "Bob"},
+				},
+			},
+		},
+		Throws: true,
+	}
+
+	assert.Equal(t, ast.Dump(expected), ast.Dump(tree.Node))
+}
+
+func TestOptimize_filter_first(t *testing.T) {
+	tree, err := parser.Parse(`first(filter(users, .Name == "Bob"))`)
+	require.NoError(t, err)
+
+	err = optimizer.Optimize(&tree.Node, nil)
+	require.NoError(t, err)
+
+	expected := &ast.BuiltinNode{
+		Name: "find",
+		Arguments: []ast.Node{
+			&ast.IdentifierNode{Value: "users"},
+			&ast.ClosureNode{
+				Node: &ast.BinaryNode{
+					Operator: "==",
+					Left: &ast.MemberNode{
+						Node:     &ast.PointerNode{},
+						Property: &ast.StringNode{Value: "Name"},
+					},
+					Right: &ast.StringNode{Value: "Bob"},
+				},
+			},
+		},
+		Throws: false,
+	}
+
+	assert.Equal(t, ast.Dump(expected), ast.Dump(tree.Node))
+}

vm/opcodes.go

@@ -71,6 +71,7 @@ const (
 	OpGetCount
 	OpGetLen
 	OpPointer
+	OpThrow
 	OpBegin
 	OpEnd // This opcode must be at the end of this list.
 )

vm/program.go

@@ -286,6 +286,9 @@ func (program *Program) Opcodes(w io.Writer) {
 		case OpPointer:
 			code("OpPointer")
 
+		case OpThrow:
+			code("OpThrow")
+
 		case OpBegin:
 			code("OpBegin")
 

vm/vm.go

@@ -457,6 +457,9 @@ func (vm *VM) Run(program *Program, env interface{}) (_ interface{}, err error)
 			scope := vm.Scope()
 			vm.push(scope.Array.Index(scope.Index).Interface())
 
+		case OpThrow:
+			panic(vm.pop().(error))
+
 		case OpBegin:
 			a := vm.pop()
 			array := reflect.ValueOf(a)