Do not optimize filter_map if #index used

Fixes #758

Author: antonmedv

expr_test.go

@@ -2763,3 +2763,24 @@ func TestExpr_env_types_map_error(t *testing.T) {
 	_, err = expr.Run(program, envTypes)
 	require.Error(t, err)
 }
+
+func TestIssue758_FilterMapIndex(t *testing.T) {
+	env := map[string]interface{}{}
+
+	exprStr := `
+        let a_map = 0..5 | filter(# % 2 == 0) | map(#index);
+        let b_filter = 0..5 | filter(# % 2 == 0);
+        let b_map = b_filter | map(#index);
+        [a_map, b_map]
+    `
+
+	result, err := expr.Eval(exprStr, env)
+	require.NoError(t, err)
+
+	expected := []interface{}{
+		[]interface{}{0, 1, 2},
+		[]interface{}{0, 1, 2},
+	}
+
+	require.Equal(t, expected, result)
+}

optimizer/filter_map.go

@@ -9,7 +9,8 @@ type filterMap struct{}
 func (*filterMap) Visit(node *Node) {
 	if mapBuiltin, ok := (*node).(*BuiltinNode); ok &&
 		mapBuiltin.Name == "map" &&
-		len(mapBuiltin.Arguments) == 2 {
+		len(mapBuiltin.Arguments) == 2 &&
+		Find(mapBuiltin.Arguments[1], isIndexPointer) == nil {
 		if predicate, ok := mapBuiltin.Arguments[1].(*PredicateNode); ok {
 			if filter, ok := mapBuiltin.Arguments[0].(*BuiltinNode); ok &&
 				filter.Name == "filter" &&
@@ -23,3 +24,10 @@ func (*filterMap) Visit(node *Node) {
 		}
 	}
 }
+
+func isIndexPointer(node Node) bool {
+	if pointer, ok := node.(*PointerNode); ok && pointer.Name == "index" {
+		return true
+	}
+	return false
+}

optimizer/filter_map_test.go

@@ -0,0 +1,74 @@
+package optimizer_test
+
+import (
+	"testing"
+
+	. "github.com/expr-lang/expr/ast"
+	"github.com/expr-lang/expr/internal/testify/assert"
+	"github.com/expr-lang/expr/internal/testify/require"
+	"github.com/expr-lang/expr/optimizer"
+	"github.com/expr-lang/expr/parser"
+)
+
+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 := &BuiltinNode{
+		Name: "filter",
+		Arguments: []Node{
+			&IdentifierNode{Value: "users"},
+			&PredicateNode{
+				Node: &BinaryNode{
+					Operator: "==",
+					Left: &MemberNode{
+						Node:     &PointerNode{},
+						Property: &StringNode{Value: "Name"},
+					},
+					Right: &StringNode{Value: "Bob"},
+				},
+			},
+		},
+		Map: &MemberNode{
+			Node:     &PointerNode{},
+			Property: &StringNode{Value: "Age"},
+		},
+	}
+
+	assert.Equal(t, Dump(expected), Dump(tree.Node))
+}
+
+func TestOptimize_filter_map_with_index_pointer(t *testing.T) {
+	tree, err := parser.Parse(`map(filter(users, true), #index)`)
+	require.NoError(t, err)
+
+	err = optimizer.Optimize(&tree.Node, nil)
+	require.NoError(t, err)
+
+	expected := &BuiltinNode{
+		Name: "map",
+		Arguments: []Node{
+			&BuiltinNode{
+				Name: "filter",
+				Arguments: []Node{
+					&IdentifierNode{Value: "users"},
+					&PredicateNode{
+						Node: &BoolNode{Value: true},
+					},
+				},
+				Throws: false,
+				Map:    nil,
+			},
+			&PredicateNode{
+				Node: &PointerNode{Name: "index"},
+			},
+		},
+		Throws: false,
+		Map:    nil,
+	}
+
+	assert.Equal(t, Dump(expected), Dump(tree.Node))
+}

optimizer/optimizer_test.go

@@ -312,37 +312,6 @@ 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.PredicateNode{
-				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)