Add interface support to operator overloading

Author: antonmedv

checker/checker.go

@@ -174,18 +174,9 @@ func (v *visitor) BinaryNode(node *ast.BinaryNode) reflect.Type {
 
 	// check operator overloading
 	if fns, ok := v.operators[node.Operator]; ok {
-		for _, fn := range fns {
-			fnType := v.types[fn]
-			firstInIndex := 0
-			if fnType.Method {
-				firstInIndex = 1 // As first argument to method is receiver.
-			}
-			firstArgType := fnType.Type.In(firstInIndex)
-			secondArgType := fnType.Type.In(firstInIndex + 1)
-
-			if l == firstArgType && r == secondArgType {
-				return fnType.Type.Out(0)
-			}
+		t, _, ok := conf.FindSuitableOperatorOverload(fns, v.types, l, r)
+		if ok {
+			return t
 		}
 	}
 

checker/patcher.go

@@ -25,24 +25,16 @@ func (p *operatorPatcher) Exit(node *ast.Node) {
 
 	leftType := binaryNode.Left.GetType()
 	rightType := binaryNode.Right.GetType()
-	for _, fn := range fns {
-		fnType := p.types[fn]
-		firstInIndex := 0
-		if fnType.Method {
-			firstInIndex = 1 // As first argument to method is receiver.
-		}
-		firstArgType := fnType.Type.In(firstInIndex)
-		secondArgType := fnType.Type.In(firstInIndex + 1)
 
-		if leftType == firstArgType && rightType == secondArgType {
-			newNode := &ast.FunctionNode{
-				Name:      fn,
-				Arguments: []ast.Node{binaryNode.Left, binaryNode.Right},
-			}
-			newNode.SetType((*node).GetType())
-			newNode.SetLocation((*node).GetLocation())
-			*node = newNode
+	_, fn, ok := conf.FindSuitableOperatorOverload(fns, p.types, leftType, rightType)
+	if ok {
+		newNode := &ast.FunctionNode{
+			Name:      fn,
+			Arguments: []ast.Node{binaryNode.Left, binaryNode.Right},
 		}
+		newNode.SetType((*node).GetType())
+		newNode.SetLocation((*node).GetLocation())
+		*node = newNode
 	}
 }
 

expr_test.go

@@ -416,6 +416,26 @@ func TestOperator_struct(t *testing.T) {
 	require.Equal(t, true, output)
 }
 
+func TestOperator_interface(t *testing.T) {
+	env := &mockEnv{
+		Ticket: &ticket{Price: 100},
+	}
+
+	code := `Ticket == "$100" && "$100" == Ticket && Now != Ticket && Now == Now`
+
+	program, err := expr.Compile(
+		code,
+		expr.Env(&mockEnv{}),
+		expr.Operator("==", "StringerStringEqual", "StringStringerEqual", "StringerStringerEqual"),
+		expr.Operator("!=", "NotStringerStringEqual", "NotStringStringerEqual", "NotStringerStringerEqual"),
+	)
+	require.NoError(t, err)
+
+	output, err := expr.Run(program, env)
+	require.NoError(t, err)
+	require.Equal(t, true, output)
+}
+
 func TestExpr(t *testing.T) {
 	env := &mockEnv{
 		Any:     "any",
@@ -836,6 +856,30 @@ func (*mockEnv) DateEqual(date time.Time, s string) bool {
 	return date.Format("2006-01-02") == s
 }
 
+func (*mockEnv) StringerStringEqual(f fmt.Stringer, s string) bool {
+	return f.String() == s
+}
+
+func (*mockEnv) StringStringerEqual(s string, f fmt.Stringer) bool {
+	return s == f.String()
+}
+
+func (*mockEnv) StringerStringerEqual(f fmt.Stringer, g fmt.Stringer) bool {
+	return f.String() == g.String()
+}
+
+func (*mockEnv) NotStringerStringEqual(f fmt.Stringer, s string) bool {
+	return f.String() != s
+}
+
+func (*mockEnv) NotStringStringerEqual(s string, f fmt.Stringer) bool {
+	return s != f.String()
+}
+
+func (*mockEnv) NotStringerStringerEqual(f fmt.Stringer, g fmt.Stringer) bool {
+	return f.String() != g.String()
+}
+
 type ticket struct {
 	Price int
 }

internal/conf/operators_table.go

@@ -1,5 +1,26 @@
 package conf
 
+import "reflect"
+
 // 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
+
+func FindSuitableOperatorOverload(fns []string, types TypesTable, l, r reflect.Type) (reflect.Type, string, bool) {
+	for _, fn := range fns {
+		fnType := types[fn]
+		firstInIndex := 0
+		if fnType.Method {
+			firstInIndex = 1 // As first argument to method is receiver.
+		}
+		firstArgType := fnType.Type.In(firstInIndex)
+		secondArgType := fnType.Type.In(firstInIndex + 1)
+
+		firstArgumentFit := l == firstArgType || (firstArgType.Kind() == reflect.Interface && l.Implements(firstArgType))
+		secondArgumentFit := r == secondArgType || (secondArgType.Kind() == reflect.Interface && r.Implements(secondArgType))
+		if firstArgumentFit && secondArgumentFit {
+			return fnType.Type.Out(0), fn, true
+		}
+	}
+	return nil, "", false
+}