Add var in a..b optimization

Author: antonmedv

expr.go

@@ -80,7 +80,10 @@ func Optimize(b bool) conf.Option {
 
 // Compile parses and compiles given input expression to bytecode program.
 func Compile(input string, ops ...conf.Option) (*vm.Program, error) {
-	config := &conf.Config{Operators: make(map[string][]string)}
+	config := &conf.Config{
+		Operators: make(map[string][]string),
+		Optimize:  true,
+	}
 
 	for _, op := range ops {
 		op(config)
@@ -103,7 +106,9 @@ func Compile(input string, ops ...conf.Option) (*vm.Program, error) {
 		checker.PatchOperators(tree, config)
 	}
 
-	optimizer.Optimize(&tree.Node)
+	if config.Optimize {
+		optimizer.Optimize(&tree.Node)
+	}
 
 	program, err := compiler.Compile(tree, config)
 	if err != nil {

expr_test.go

@@ -448,25 +448,21 @@ func TestExpr(t *testing.T) {
 			4,
 		},
 		{
-			`[1, 2, 3]`,
-			[]int{1, 2, 3},
+			`len([1, 2, 3])`,
+			3,
 		},
 		{
-			`[1, Two, 3]`,
-			[]interface{}{1, 2, 3},
+			`len([1, Two, 3])`,
+			3,
 		},
 		{
-			`["hello", "world"]`,
-			[]string{"hello", "world"},
+			`len(["hello", "world"])`,
+			2,
 		},
 		{
 			`{foo: 0, bar: 1}`,
 			map[string]interface{}{"foo": 0, "bar": 1},
 		},
-		{
-			`[1, 2, 3]`,
-			[]int{1, 2, 3},
-		},
 		{
 			`{foo: 0, bar: 1}`,
 			map[string]interface{}{"foo": 0, "bar": 1},
@@ -583,6 +579,10 @@ func TestExpr(t *testing.T) {
 			`Array[:] == Array`,
 			true,
 		},
+		{
+			`One in 0..1 && Two not in 0..1`,
+			true,
+		},
 	}
 
 	for _, tt := range tests {
@@ -594,6 +594,23 @@ func TestExpr(t *testing.T) {
 
 		assert.Equal(t, tt.want, got, tt.code)
 	}
+
+	for _, tt := range tests {
+		program, err := expr.Compile(tt.code, expr.Optimize(false))
+		require.NoError(t, err, tt.code)
+
+		got, err := expr.Run(program, env)
+		require.NoError(t, err, tt.code)
+
+		assert.Equal(t, tt.want, got, "unoptimized: "+tt.code)
+	}
+
+	for _, tt := range tests {
+		got, err := expr.Eval(tt.code, env)
+		require.NoError(t, err, tt.code)
+
+		assert.Equal(t, tt.want, got, "eval: "+tt.code)
+	}
 }
 
 type mockEnv struct {

optimizer/optimizer.go

@@ -1,46 +1,45 @@
 package optimizer
 
 import (
-	"github.com/antonmedv/expr/ast"
+	. "github.com/antonmedv/expr/ast"
 )
 
-type optimizer struct {
-}
-
-func (*optimizer) Enter(node *ast.Node) {}
+type fold struct{}
+type inRange struct{}
 
-func (*optimizer) Exit(node *ast.Node) {
+func (*fold) Enter(node *Node) {}
+func (*fold) Exit(node *Node) {
 	switch n := (*node).(type) {
-	case *ast.ArrayNode:
+	case *ArrayNode:
 		if len(n.Nodes) > 0 {
 
 			for _, a := range n.Nodes {
-				if _, ok := a.(*ast.IntegerNode); !ok {
+				if _, ok := a.(*IntegerNode); !ok {
 					goto string
 				}
 			}
 			{
 				value := make([]int, len(n.Nodes))
 				for i, a := range n.Nodes {
-					value[i] = a.(*ast.IntegerNode).Value
+					value[i] = a.(*IntegerNode).Value
 				}
-				*node = &ast.ConstantNode{
+				*node = &ConstantNode{
 					Value: value,
 				}
 			}
 
 		string:
 			for _, a := range n.Nodes {
-				if _, ok := a.(*ast.StringNode); !ok {
+				if _, ok := a.(*StringNode); !ok {
 					return
 				}
 			}
 			{
 				value := make([]string, len(n.Nodes))
 				for i, a := range n.Nodes {
-					value[i] = a.(*ast.StringNode).Value
+					value[i] = a.(*StringNode).Value
 				}
-				*node = &ast.ConstantNode{
+				*node = &ConstantNode{
 					Value: value,
 				}
 			}
@@ -49,7 +48,41 @@ func (*optimizer) Exit(node *ast.Node) {
 	}
 }
 
-func Optimize(node *ast.Node) {
-	o := &optimizer{}
-	ast.Walk(node, o)
+func (*inRange) Enter(node *Node) {}
+func (*inRange) Exit(node *Node) {
+	switch n := (*node).(type) {
+	case *BinaryNode:
+		if n.Operator == "in" || n.Operator == "not in" {
+			if rng, ok := n.Right.(*BinaryNode); ok && rng.Operator == ".." {
+				if from, ok := n.Left.(*IntegerNode); ok {
+					if to, ok := n.Right.(*IntegerNode); ok {
+						*node = &BinaryNode{
+							Operator: "and",
+							Left: &BinaryNode{
+								Operator: ">=",
+								Left:     n.Left,
+								Right:    from,
+							},
+							Right: &BinaryNode{
+								Operator: "<=",
+								Left:     n.Left,
+								Right:    to,
+							},
+						}
+						if n.Operator == "not in" {
+							*node = &UnaryNode{
+								Operator: "not",
+								Node:     *node,
+							}
+						}
+					}
+				}
+			}
+		}
+	}
+}
+
+func Optimize(node *Node) {
+	Walk(node, &fold{})
+	Walk(node, &inRange{})
 }

optimizer/optimizer_test.go

@@ -24,3 +24,28 @@ func TestOptimize_constant_folding(t *testing.T) {
 
 	assert.Equal(t, litter.Sdump(expected), litter.Sdump(tree.Node))
 }
+
+func TestOptimize_in_range(t *testing.T) {
+	tree, err := parser.Parse(`age in 18..31`)
+	require.NoError(t, err)
+
+	optimizer.Optimize(&tree.Node)
+
+	expected := &ast.BinaryNode{
+		Operator: "in",
+		Left: &ast.IdentifierNode{
+			Value: "age",
+		},
+		Right: &ast.BinaryNode{
+			Operator: "..",
+			Left: &ast.IntegerNode{
+				Value: 18,
+			},
+			Right: &ast.IntegerNode{
+				Value: 31,
+			},
+		},
+	}
+
+	assert.Equal(t, litter.Sdump(expected), litter.Sdump(tree.Node))
+}