Add map(filter()) optimization

Author: antonmedv

ast/node.go

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

compiler/compiler.go

@@ -688,7 +688,11 @@ func (c *compiler) BuiltinNode(node *ast.BuiltinNode) {
 			c.compile(node.Arguments[1])
 			c.emitCond(func() {
 				c.emit(OpIncrementCount)
-				c.emit(OpPointer)
+				if node.Map != nil {
+					c.compile(node.Map)
+				} else {
+					c.emit(OpPointer)
+				}
 			})
 		})
 		c.emit(OpGetCount)
@@ -728,7 +732,11 @@ func (c *compiler) BuiltinNode(node *ast.BuiltinNode) {
 			c.compile(node.Arguments[1])
 			noop := c.emit(OpJumpIfFalse, placeholder)
 			c.emit(OpPop)
-			c.emit(OpPointer)
+			if node.Map != nil {
+				c.compile(node.Map)
+			} else {
+				c.emit(OpPointer)
+			}
 			loopBreak = c.emit(OpJump, placeholder)
 			c.patchJump(noop)
 			c.emit(OpPop)
@@ -769,7 +777,11 @@ func (c *compiler) BuiltinNode(node *ast.BuiltinNode) {
 			c.compile(node.Arguments[1])
 			noop := c.emit(OpJumpIfFalse, placeholder)
 			c.emit(OpPop)
-			c.emit(OpPointer)
+			if node.Map != nil {
+				c.compile(node.Map)
+			} else {
+				c.emit(OpPointer)
+			}
 			loopBreak = c.emit(OpJump, placeholder)
 			c.patchJump(noop)
 			c.emit(OpPop)

expr_test.go

@@ -1000,6 +1000,34 @@ func TestExpr(t *testing.T) {
 			`last(filter(1..9, # % 2 == 0))`,
 			8,
 		},
+		{
+			`map(filter(1..9, # % 2 == 0), # * 2)`,
+			[]interface{}{4, 8, 12, 16},
+		},
+		{
+			`map(map(filter(1..9, # % 2 == 0), # * 2), # * 2)`,
+			[]interface{}{8, 16, 24, 32},
+		},
+		{
+			`first(map(filter(1..9, # % 2 == 0), # * 2))`,
+			4,
+		},
+		{
+			`map(filter(1..9, # % 2 == 0), # * 2)[-1]`,
+			16,
+		},
+		{
+			`len(map(filter(1..9, # % 2 == 0), # * 2))`,
+			4,
+		},
+		{
+			`len(filter(map(1..9, # * 2), # % 2 == 0))`,
+			9,
+		},
+		{
+			`first(filter(map(1..9, # * 2), # % 2 == 0))`,
+			2,
+		},
 	}
 
 	for _, tt := range tests {

optimizer/filter_first.go

@@ -16,6 +16,7 @@ func (*filterFirst) Visit(node *Node) {
 					Name:      "find",
 					Arguments: filter.Arguments,
 					Throws:    true, // to match the behavior of filter()[0]
+					Map:       filter.Map,
 				})
 			}
 		}
@@ -30,6 +31,7 @@ func (*filterFirst) Visit(node *Node) {
 				Name:      "find",
 				Arguments: filter.Arguments,
 				Throws:    false, // as first() will return nil if not found
+				Map:       filter.Map,
 			})
 		}
 	}

optimizer/filter_last.go

@@ -16,6 +16,7 @@ func (*filterLast) Visit(node *Node) {
 					Name:      "findLast",
 					Arguments: filter.Arguments,
 					Throws:    true, // to match the behavior of filter()[-1]
+					Map:       filter.Map,
 				})
 			}
 		}
@@ -30,6 +31,7 @@ func (*filterLast) Visit(node *Node) {
 				Name:      "findLast",
 				Arguments: filter.Arguments,
 				Throws:    false, // as last() will return nil if not found
+				Map:       filter.Map,
 			})
 		}
 	}

optimizer/filter_map.go

@@ -0,0 +1,25 @@
+package optimizer
+
+import (
+	. "github.com/antonmedv/expr/ast"
+)
+
+type filterMap struct{}
+
+func (*filterMap) Visit(node *Node) {
+	if mapBuiltin, ok := (*node).(*BuiltinNode); ok &&
+		mapBuiltin.Name == "map" &&
+		len(mapBuiltin.Arguments) == 2 {
+		if closure, ok := mapBuiltin.Arguments[1].(*ClosureNode); ok {
+			if filter, ok := mapBuiltin.Arguments[0].(*BuiltinNode); ok &&
+				filter.Name == "filter" &&
+				filter.Map == nil /* not already optimized */ {
+				Patch(node, &BuiltinNode{
+					Name:      "filter",
+					Arguments: filter.Arguments,
+					Map:       closure.Node,
+				})
+			}
+		}
+	}
+}

optimizer/optimizer.go

@@ -33,8 +33,9 @@ func Optimize(node *Node, config *conf.Config) error {
 	}
 	Walk(node, &inRange{})
 	Walk(node, &constRange{})
+	Walk(node, &filterMap{})
 	Walk(node, &filterLen{})
-	Walk(node, &filterFirst{})
 	Walk(node, &filterLast{})
+	Walk(node, &filterFirst{})
 	return nil
 }

optimizer/optimizer_test.go

@@ -277,3 +277,66 @@ func TestOptimize_filter_last(t *testing.T) {
 
 	assert.Equal(t, ast.Dump(expected), ast.Dump(tree.Node))
 }
+
+func TestOptimize_filter_map(t *testing.T) {
+	tree, err := parser.Parse(`map(filter(users, .Name == "Bob"), .Age)`)
+	require.NoError(t, err)
+
+	err = optimizer.Optimize(&tree.Node, nil)
+	require.NoError(t, err)
+
+	expected := &ast.BuiltinNode{
+		Name: "filter",
+		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"},
+				},
+			},
+		},
+		Map: &ast.MemberNode{
+			Node:     &ast.PointerNode{},
+			Property: &ast.StringNode{Value: "Age"},
+		},
+	}
+
+	assert.Equal(t, ast.Dump(expected), ast.Dump(tree.Node))
+}
+
+func TestOptimize_filter_map_first(t *testing.T) {
+	tree, err := parser.Parse(`first(map(filter(users, .Name == "Bob"), .Age))`)
+	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"},
+				},
+			},
+		},
+		Map: &ast.MemberNode{
+			Node:     &ast.PointerNode{},
+			Property: &ast.StringNode{Value: "Age"},
+		},
+		Throws: false,
+	}
+
+	assert.Equal(t, ast.Dump(expected), ast.Dump(tree.Node))
+}