Add in array optimization

Author: antonmedv

compiler/compiler.go

@@ -74,7 +74,7 @@ func (c *compiler) emitPush(value interface{}) int {
 func (c *compiler) makeConstant(i interface{}) []byte {
 	hashable := true
 	switch reflect.TypeOf(i).Kind() {
-	case reflect.Slice:
+	case reflect.Slice, reflect.Map:
 		hashable = false
 	}
 

optimizer/optimizer.go

@@ -5,16 +5,75 @@ import (
 	"math"
 )
 
+type inArray struct{}
 type fold struct {
 	applied bool
 }
 type inRange struct{}
 type constRange struct{}
 
-func patch(node *Node, newNode Node) {
-	newNode.SetType((*node).GetType())
-	newNode.SetLocation((*node).GetLocation())
-	*node = newNode
+func Optimize(node *Node) {
+	Walk(node, &inArray{})
+	limit := 1000
+	for {
+		fold := &fold{}
+		Walk(node, fold)
+		limit--
+		if !fold.applied || limit == 0 {
+			break
+		}
+	}
+	Walk(node, &inRange{})
+	Walk(node, &constRange{})
+}
+
+func (*inArray) Enter(node *Node) {}
+func (*inArray) Exit(node *Node) {
+	switch n := (*node).(type) {
+	case *BinaryNode:
+		if n.Operator == "in" || n.Operator == "not in" {
+			if array, ok := n.Right.(*ArrayNode); ok {
+				if len(array.Nodes) > 0 {
+
+					for _, a := range array.Nodes {
+						if _, ok := a.(*IntegerNode); !ok {
+							goto string
+						}
+					}
+					{
+						value := make(map[int]bool)
+						for _, a := range array.Nodes {
+							value[a.(*IntegerNode).Value] = true
+						}
+						patch(node, &BinaryNode{
+							Operator: n.Operator,
+							Left:     n.Left,
+							Right:    &ConstantNode{Value: value},
+						})
+					}
+
+				string:
+					for _, a := range array.Nodes {
+						if _, ok := a.(*StringNode); !ok {
+							return
+						}
+					}
+					{
+						value := make(map[string]bool)
+						for _, a := range array.Nodes {
+							value[a.(*StringNode).Value] = true
+						}
+						patch(node, &BinaryNode{
+							Operator: n.Operator,
+							Left:     n.Left,
+							Right:    &ConstantNode{Value: value},
+						})
+					}
+
+				}
+			}
+		}
+	}
 }
 
 func (*fold) Enter(node *Node) {}
@@ -171,16 +230,8 @@ func (*constRange) Exit(node *Node) {
 	}
 }
 
-func Optimize(node *Node) {
-	limit := 1000
-	for {
-		fold := &fold{}
-		Walk(node, fold)
-		limit--
-		if !fold.applied || limit == 0 {
-			break
-		}
-	}
-	Walk(node, &inRange{})
-	Walk(node, &constRange{})
+func patch(node *Node, newNode Node) {
+	newNode.SetType((*node).GetType())
+	newNode.SetLocation((*node).GetLocation())
+	*node = newNode
 }

optimizer/optimizer_test.go

@@ -11,15 +11,29 @@ import (
 )
 
 func TestOptimize_constant_folding(t *testing.T) {
-	tree, err := parser.Parse(`v in [1,2,3,4,5]`)
+	tree, err := parser.Parse(`[1,2,3][5*5-25]`)
+	require.NoError(t, err)
+
+	optimizer.Optimize(&tree.Node)
+
+	expected := &ast.IndexNode{
+		Node:  &ast.ConstantNode{Value: []int{1, 2, 3}},
+		Index: &ast.IntegerNode{Value: 0},
+	}
+
+	assert.Equal(t, litter.Sdump(expected), litter.Sdump(tree.Node))
+}
+
+func TestOptimize_in_array(t *testing.T) {
+	tree, err := parser.Parse(`v in [1,2,3]`)
 	require.NoError(t, err)
 
 	optimizer.Optimize(&tree.Node)
 
 	expected := &ast.BinaryNode{
 		Operator: "in",
 		Left:     &ast.IdentifierNode{Value: "v"},
-		Right:    &ast.ConstantNode{Value: []int{1, 2, 3, 4, 5}},
+		Right:    &ast.ConstantNode{Value: map[int]bool{1: true, 2: true, 3: true}},
 	}
 
 	assert.Equal(t, litter.Sdump(expected), litter.Sdump(tree.Node))