Refactor operator override to be purely patcher pkg

Author: antonmedv

conf/config.go

@@ -9,34 +9,28 @@ import (
 	"github.com/expr-lang/expr/vm/runtime"
 )
 
-type FunctionTable map[string]*builtin.Function
-
-// OperatorsTable maps binary operators to corresponding list of functions.
-// Functions should be provided in the environment to allow operator overloading.
-type OperatorsTable map[string][]string
+type FunctionsTable map[string]*builtin.Function
 
 type Config struct {
 	Env         any
 	Types       TypesTable
 	MapEnv      bool
 	DefaultType reflect.Type
-	Operators   OperatorsTable
 	Expect      reflect.Kind
 	ExpectAny   bool
 	Optimize    bool
 	Strict      bool
 	ConstFns    map[string]reflect.Value
 	Visitors    []ast.Visitor
-	Functions   map[string]*builtin.Function
-	Builtins    map[string]*builtin.Function
+	Functions   FunctionsTable
+	Builtins    FunctionsTable
 	Disabled    map[string]bool // disabled builtins
 }
 
 // CreateNew creates new config with default values.
 func CreateNew() *Config {
 	c := &Config{
 		Optimize:  true,
-		Operators: make(map[string][]string),
 		ConstFns:  make(map[string]reflect.Value),
 		Functions: make(map[string]*builtin.Function),
 		Builtins:  make(map[string]*builtin.Function),
@@ -73,10 +67,6 @@ func (c *Config) WithEnv(env any) {
 	c.Strict = true
 }
 
-func (c *Config) Operator(operator string, fns ...string) {
-	c.Operators[operator] = append(c.Operators[operator], fns...)
-}
-
 func (c *Config) ConstExpr(name string) {
 	if c.Env == nil {
 		panic("no environment is specified for ConstExpr()")
@@ -88,42 +78,14 @@ func (c *Config) ConstExpr(name string) {
 	c.ConstFns[name] = fn
 }
 
-func (c *Config) Check() {
-	for operator, fns := range c.Operators {
-		for _, fn := range fns {
-			fnType, foundType := c.Types[fn]
-			fnFunc, foundFunc := c.Functions[fn]
-			if !foundFunc && (!foundType || fnType.Type.Kind() != reflect.Func) {
-				panic(fmt.Errorf("function %s for %s operator does not exist in the environment", fn, operator))
-			}
-
-			if foundType {
-				checkType(fnType, fn, operator)
-			}
-			if foundFunc {
-				checkFunc(fnFunc, fn, operator)
-			}
-		}
-	}
+type Checker interface {
+	Check()
 }
 
-func checkType(fnType Tag, fn string, operator string) {
-	requiredNumIn := 2
-	if fnType.Method {
-		requiredNumIn = 3 // As first argument of method is receiver.
-	}
-	if fnType.Type.NumIn() != requiredNumIn || fnType.Type.NumOut() != 1 {
-		panic(fmt.Errorf("function %s for %s operator does not have a correct signature", fn, operator))
-	}
-}
-
-func checkFunc(fn *builtin.Function, name string, operator string) {
-	if len(fn.Types) == 0 {
-		panic(fmt.Errorf("function %s for %s operator misses types", name, operator))
-	}
-	for _, t := range fn.Types {
-		if t.NumIn() != 2 || t.NumOut() != 1 {
-			panic(fmt.Errorf("function %s for %s operator does not have a correct signature", name, operator))
+func (c *Config) Check() {
+	for _, v := range c.Visitors {
+		if c, ok := v.(Checker); ok {
+			c.Check()
 		}
 	}
 }

conf/functions.go

@@ -42,7 +42,13 @@ func AllowUndefinedVariables() Option {
 // Operator allows to replace a binary operator with a function.
 func Operator(operator string, fn ...string) Option {
 	return func(c *conf.Config) {
-		c.Operator(operator, fn...)
+		p := &patcher.OperatorOverride{
+			Operator:  operator,
+			Overrides: fn,
+			Types:     c.Types,
+			Functions: c.Functions,
+		}
+		c.Visitors = append(c.Visitors, p)
 	}
 }
 
@@ -188,14 +194,6 @@ func Compile(input string, ops ...Option) (*vm.Program, error) {
 	}
 	config.Check()
 
-	if len(config.Operators) > 0 {
-		config.Visitors = append(config.Visitors, &patcher.Operator{
-			Operators: config.Operators,
-			Types:     config.Types,
-			Functions: config.Functions,
-		})
-	}
-
 	tree, err := parser.ParseWithConfig(input, config)
 	if err != nil {
 		return nil, err

conf/operators.go

@@ -1,31 +1,35 @@
 package patcher
 
 import (
+	"fmt"
+	"reflect"
+
 	"github.com/expr-lang/expr/ast"
+	"github.com/expr-lang/expr/builtin"
 	"github.com/expr-lang/expr/conf"
 )
 
-type Operator struct {
-	Operators conf.OperatorsTable
-	Types     conf.TypesTable
-	Functions conf.FunctionTable
+type OperatorOverride struct {
+	Operator  string              // Operator token to override.
+	Overrides []string            // List of function names to override operator with.
+	Types     conf.TypesTable     // Env types.
+	Functions conf.FunctionsTable // Env functions.
 }
 
-func (p *Operator) Visit(node *ast.Node) {
+func (p *OperatorOverride) Visit(node *ast.Node) {
 	binaryNode, ok := (*node).(*ast.BinaryNode)
 	if !ok {
 		return
 	}
 
-	fns, ok := p.Operators[binaryNode.Operator]
-	if !ok {
+	if binaryNode.Operator != p.Operator {
 		return
 	}
 
 	leftType := binaryNode.Left.Type()
 	rightType := binaryNode.Right.Type()
 
-	ret, fn, ok := conf.FindSuitableOperatorOverload(fns, p.Types, p.Functions, leftType, rightType)
+	ret, fn, ok := p.FindSuitableOperatorOverload(leftType, rightType)
 	if ok {
 		newNode := &ast.CallNode{
 			Callee:    &ast.IdentifierNode{Value: fn},
@@ -35,3 +39,97 @@ func (p *Operator) Visit(node *ast.Node) {
 		ast.Patch(node, newNode)
 	}
 }
+
+func (p *OperatorOverride) FindSuitableOperatorOverload(l, r reflect.Type) (reflect.Type, string, bool) {
+	t, fn, ok := p.findSuitableOperatorOverloadInFunctions(l, r)
+	if !ok {
+		t, fn, ok = p.findSuitableOperatorOverloadInTypes(l, r)
+	}
+	return t, fn, ok
+}
+
+func (p *OperatorOverride) findSuitableOperatorOverloadInTypes(l, r reflect.Type) (reflect.Type, string, bool) {
+	for _, fn := range p.Overrides {
+		fnType, ok := p.Types[fn]
+		if !ok {
+			continue
+		}
+		firstInIndex := 0
+		if fnType.Method {
+			firstInIndex = 1 // As first argument to method is receiver.
+		}
+		ret, done := checkTypeSuits(fnType.Type, l, r, firstInIndex)
+		if done {
+			return ret, fn, true
+		}
+	}
+	return nil, "", false
+}
+
+func (p *OperatorOverride) findSuitableOperatorOverloadInFunctions(l, r reflect.Type) (reflect.Type, string, bool) {
+	for _, fn := range p.Overrides {
+		fnType, ok := p.Functions[fn]
+		if !ok {
+			continue
+		}
+		firstInIndex := 0
+		for _, overload := range fnType.Types {
+			ret, done := checkTypeSuits(overload, l, r, firstInIndex)
+			if done {
+				return ret, fn, true
+			}
+		}
+	}
+	return nil, "", false
+}
+
+func checkTypeSuits(t reflect.Type, l reflect.Type, r reflect.Type, firstInIndex int) (reflect.Type, bool) {
+	firstArgType := t.In(firstInIndex)
+	secondArgType := t.In(firstInIndex + 1)
+
+	firstArgumentFit := l == firstArgType || (firstArgType.Kind() == reflect.Interface && (l == nil || l.Implements(firstArgType)))
+	secondArgumentFit := r == secondArgType || (secondArgType.Kind() == reflect.Interface && (r == nil || r.Implements(secondArgType)))
+	if firstArgumentFit && secondArgumentFit {
+		return t.Out(0), true
+	}
+	return nil, false
+}
+
+func (p *OperatorOverride) Check() {
+	for _, fn := range p.Overrides {
+		fnType, foundType := p.Types[fn]
+		fnFunc, foundFunc := p.Functions[fn]
+		if !foundFunc && (!foundType || fnType.Type.Kind() != reflect.Func) {
+			panic(fmt.Errorf("function %s for %s operator does not exist in the environment", fn, p.Operator))
+		}
+
+		if foundType {
+			checkType(fnType, fn, p.Operator)
+		}
+
+		if foundFunc {
+			checkFunc(fnFunc, fn, p.Operator)
+		}
+	}
+}
+
+func checkType(fnType conf.Tag, fn string, operator string) {
+	requiredNumIn := 2
+	if fnType.Method {
+		requiredNumIn = 3 // As first argument of method is receiver.
+	}
+	if fnType.Type.NumIn() != requiredNumIn || fnType.Type.NumOut() != 1 {
+		panic(fmt.Errorf("function %s for %s operator does not have a correct signature", fn, operator))
+	}
+}
+
+func checkFunc(fn *builtin.Function, name string, operator string) {
+	if len(fn.Types) == 0 {
+		panic(fmt.Errorf("function %s for %s operator misses types", name, operator))
+	}
+	for _, t := range fn.Types {
+		if t.NumIn() != 2 || t.NumOut() != 1 {
+			panic(fmt.Errorf("function %s for %s operator does not have a correct signature", name, operator))
+		}
+	}
+}

expr.go

@@ -26,12 +26,6 @@ func TestOperator_struct(t *testing.T) {
 	require.Equal(t, true, output)
 }
 
-func TestOperator_options_another_order(t *testing.T) {
-	code := `Time == "2017-10-23"`
-	_, err := expr.Compile(code, expr.Operator("==", "TimeEqualString"), expr.Env(mock.Env{}))
-	require.NoError(t, err)
-}
-
 func TestOperator_no_env(t *testing.T) {
 	code := `Time == "2017-10-23"`
 	require.Panics(t, func() {