Refactor operator patcher

compiler/patcher.go

@@ -1,44 +0,0 @@
-package compiler
-
-import (
-	"github.com/antonmedv/expr/ast"
-	"github.com/antonmedv/expr/conf"
-)
-
-type operatorPatcher struct {
-	ops   map[string][]string
-	types conf.TypesTable
-}
-
-func (p *operatorPatcher) Enter(node *ast.Node) {}
-func (p *operatorPatcher) Exit(node *ast.Node) {
-	binaryNode, ok := (*node).(*ast.BinaryNode)
-	if !ok {
-		return
-	}
-
-	fns, ok := p.ops[binaryNode.Operator]
-	if !ok {
-		return
-	}
-
-	leftType := binaryNode.Left.Type()
-	rightType := binaryNode.Right.Type()
-
-	_, fn, ok := conf.FindSuitableOperatorOverload(fns, p.types, leftType, rightType)
-	if ok {
-		newNode := &ast.CallNode{
-			Callee:    &ast.IdentifierNode{Value: fn},
-			Arguments: []ast.Node{binaryNode.Left, binaryNode.Right},
-		}
-		ast.Patch(node, newNode)
-	}
-}
-
-func PatchOperators(node *ast.Node, config *conf.Config) {
-	if len(config.Operators) == 0 {
-		return
-	}
-	patcher := &operatorPatcher{ops: config.Operators, types: config.Types}
-	ast.Walk(node, patcher)
-}

conf/operators.go

@@ -1,6 +1,9 @@
 package conf
 
-import "reflect"
+import (
+	"github.com/antonmedv/expr/ast"
+	"reflect"
+)
 
 // OperatorsTable maps binary operators to corresponding list of functions.
 // Functions should be provided in the environment to allow operator overloading.
@@ -24,3 +27,33 @@ func FindSuitableOperatorOverload(fns []string, types TypesTable, l, r reflect.T
 	}
 	return nil, "", false
 }
+
+type OperatorPatcher struct {
+	Operators OperatorsTable
+	Types     TypesTable
+}
+
+func (p *OperatorPatcher) Enter(_ *ast.Node) {}
+func (p *OperatorPatcher) Exit(node *ast.Node) {
+	binaryNode, ok := (*node).(*ast.BinaryNode)
+	if !ok {
+		return
+	}
+
+	fns, ok := p.Operators[binaryNode.Operator]
+	if !ok {
+		return
+	}
+
+	leftType := binaryNode.Left.Type()
+	rightType := binaryNode.Right.Type()
+
+	_, fn, ok := FindSuitableOperatorOverload(fns, p.Types, leftType, rightType)
+	if ok {
+		newNode := &ast.CallNode{
+			Callee:    &ast.IdentifierNode{Value: fn},
+			Arguments: []ast.Node{binaryNode.Left, binaryNode.Right},
+		}
+		ast.Patch(node, newNode)
+	}
+}

expr.go

@@ -133,6 +133,13 @@ func Compile(input string, ops ...Option) (*vm.Program, error) {
 		op(config)
 	}
 
+	if len(config.Operators) > 0 {
+		config.Visitors = append(config.Visitors, &conf.OperatorPatcher{
+			Operators: config.Operators,
+			Types:     config.Types,
+		})
+	}
+
 	if err := config.Check(); err != nil {
 		return nil, err
 	}
@@ -142,25 +149,22 @@ func Compile(input string, ops ...Option) (*vm.Program, error) {
 		return nil, err
 	}
 
-	_, err = checker.Check(tree, config)
-
-	// If we have a patch to apply, it may fix out error and
-	// second type check is needed. Otherwise, it is an error.
-	if err != nil && len(config.Visitors) == 0 {
-		return nil, err
-	}
-
-	// Patch operators before Optimize, as we may also mark it as ConstExpr.
-	compiler.PatchOperators(&tree.Node, config)
-
-	if len(config.Visitors) >= 0 {
+	if len(config.Visitors) > 0 {
 		for _, v := range config.Visitors {
+			// We need to perform types check, because some visitors may rely on
+			// types information available in the tree.
+			_, _ = checker.Check(tree, config)
 			ast.Walk(&tree.Node, v)
 		}
 		_, err = checker.Check(tree, config)
 		if err != nil {
 			return nil, err
 		}
+	} else {
+		_, err = checker.Check(tree, config)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	if config.Optimize {

expr_test.go

@@ -1277,7 +1277,7 @@ func TestCompile_exposed_error(t *testing.T) {
 	require.Equal(t, `{"Line":1,"Column":2,"Message":"invalid operation: == (mismatched types int and bool)","Snippet":"\n | 1 == true\n | ..^"}`, string(b))
 }
 
-func TestAsBool_exposed_error_(t *testing.T) {
+func TestAsBool_exposed_error(t *testing.T) {
 	_, err := expr.Compile(`42`, expr.AsBool())
 	require.Error(t, err)