Allow overloading binary node with a function from env

Author: davidklassen

checker/checker.go

@@ -25,10 +25,12 @@ func Check(tree *parser.Tree, config *conf.Config) (t reflect.Type, err error) {
 	}
 	if config != nil {
 		v.types = config.Types
+		v.operators = config.Operators
 		v.expect = config.Expect
 	}
 
 	t = v.visit(tree.Node)
+	patchOperators(tree, config)
 
 	if v.expect != reflect.Invalid {
 		switch v.expect {
@@ -50,6 +52,7 @@ okay:
 
 type visitor struct {
 	types       conf.TypesTable
+	operators   conf.OperatorsTable
 	expect      reflect.Kind
 	collections []reflect.Type
 }
@@ -165,6 +168,20 @@ func (v *visitor) BinaryNode(node *ast.BinaryNode) reflect.Type {
 	l := v.visit(node.Left)
 	r := v.visit(node.Right)
 
+	// check operator overloading
+	if fns, ok := v.operators[node.Operator]; ok {
+		for _, fn := range fns {
+			fnType := v.types[fn]
+
+			firstArgType := fnType.Type.In(0)
+			secondArgType := fnType.Type.In(1)
+
+			if l == firstArgType && r == secondArgType {
+				return fnType.Type.Out(0)
+			}
+		}
+	}
+
 	switch node.Operator {
 	case "==", "!=":
 		if isNumber(l) && isNumber(r) {

checker/patcher.go

@@ -0,0 +1,49 @@
+package checker
+
+import (
+	"github.com/antonmedv/expr/ast"
+	"github.com/antonmedv/expr/internal/conf"
+	"github.com/antonmedv/expr/parser"
+)
+
+type operatorPatcher struct {
+	ast.BaseVisitor
+	ops   map[string][]string
+	types conf.TypesTable
+}
+
+func (p *operatorPatcher) Node(node *ast.Node) {
+	binaryNode, ok := (*node).(*ast.BinaryNode)
+	if !ok {
+		return
+	}
+
+	fns, ok := p.ops[binaryNode.Operator]
+	if !ok {
+		return
+	}
+
+	leftType := binaryNode.Left.GetType()
+	rightType := binaryNode.Right.GetType()
+	for _, fn := range fns {
+		fnType := p.types[fn]
+
+		firstArgType := fnType.Type.In(0)
+		secondArgType := fnType.Type.In(1)
+
+		if leftType == firstArgType && rightType == secondArgType {
+			*node = &ast.FunctionNode{
+				Name:      fn,
+				Arguments: []ast.Node{binaryNode.Left, binaryNode.Right},
+			}
+		}
+	}
+}
+
+func patchOperators(tree *parser.Tree, config *conf.Config) {
+	if len(config.Operators) == 0 {
+		return
+	}
+	patcher := &operatorPatcher{ops: config.Operators, types: config.Types}
+	ast.Walk(&tree.Node, patcher)
+}

expr.go

@@ -43,6 +43,12 @@ func Env(i interface{}) conf.Option {
 	}
 }
 
+func Operator(operator string, fns ...string) conf.Option {
+	return func(c *conf.Config) {
+		c.Operators[operator] = append(c.Operators[operator], fns...)
+	}
+}
+
 // AsBool tells the compiler to expect boolean result.
 func AsBool() conf.Option {
 	return func(c *conf.Config) {
@@ -66,12 +72,16 @@ func AsFloat64() 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{}
+	config := &conf.Config{Operators: make(map[string][]string)}
 
 	for _, op := range ops {
 		op(config)
 	}
 
+	if err := config.Check(); err != nil {
+		return nil, err
+	}
+
 	tree, err := parser.Parse(input)
 	if err != nil {
 		return nil, err

expr_test.go

@@ -5,7 +5,10 @@ import (
 	"fmt"
 	"github.com/antonmedv/expr"
 	"github.com/antonmedv/expr/vm"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"strings"
+	"testing"
 )
 
 func ExampleEval() {
@@ -183,6 +186,49 @@ func ExampleEnv() {
 	// Output: true
 }
 
+func ExampleOperator() {
+	type Place struct {
+		Code string
+	}
+	type Segment struct {
+		Origin Place
+	}
+	type Helpers struct {
+		PlaceEq func(p Place, s string) bool
+	}
+	type Request struct {
+		Segments []*Segment
+		Helpers
+	}
+
+	code := `Segments[0].Origin == "MOW" && PlaceEq(Segments[0].Origin, "MOW")`
+
+	program, err := expr.Compile(code, expr.Env(&Request{}), expr.Operator("==", "PlaceEq"))
+	if err != nil {
+		fmt.Printf("%v", err)
+		return
+	}
+
+	request := &Request{
+		Segments: []*Segment{
+			{Origin: Place{Code: "MOW"}},
+		},
+		Helpers: Helpers{PlaceEq: func(p Place, s string) bool {
+			return p.Code == s
+		}},
+	}
+
+	output, err := expr.Run(program, request)
+	if err != nil {
+		fmt.Printf("%v", err)
+		return
+	}
+
+	fmt.Printf("%v", output)
+
+	// Output: true
+}
+
 func ExampleEval_marshal() {
 	env := map[string]int{
 		"foo": 1,
@@ -218,3 +264,93 @@ func ExampleEval_marshal() {
 
 	// Output: 3
 }
+
+func TestExpr(t *testing.T) {
+	type mockEnv struct {
+		One, Two, Three int
+		IntArray        []int
+		MultiDimArray   [][]int
+		Sum             func(list []int) int
+		Inc             func(int) int
+	}
+
+	request := mockEnv{
+		One:           1,
+		Two:           2,
+		Three:         3,
+		IntArray:      []int{1, 2, 3},
+		MultiDimArray: [][]int{{1, 2, 3}, {1, 2, 3}},
+		Sum: func(list []int) int {
+			var ret int
+			for _, el := range list {
+				ret += el
+			}
+			return ret
+		},
+		Inc: func(a int) int { return a + 1 },
+	}
+
+	tests := []struct {
+		name string
+		code string
+		want interface{}
+	}{
+		{
+			name: "+ operator",
+			code: "1 + 1",
+			want: 2,
+		},
+		{
+			name: "associativity",
+			code: "(One * Two) * Three == One * (Two * Three)",
+			want: true,
+		},
+		{
+			name: "indexing",
+			code: "IntArray[0]",
+			want: 1,
+		},
+		{
+			name: "helpers",
+			code: "Sum(IntArray)",
+			want: 6,
+		},
+		{
+			name: "binary with indexing",
+			code: "IntArray[0] < IntArray[1]",
+			want: true,
+		},
+		{
+			name: "helpers with indexing",
+			code: "Sum(MultiDimArray[0])",
+			want: 6,
+		},
+		{
+			name: "helpers with indexing in binary operations",
+			code: "Sum(MultiDimArray[0]) + Sum(MultiDimArray[1])",
+			want: 12,
+		},
+		{
+			name: "binary operations in function call arguments",
+			code: "Inc(IntArray[0] + IntArray[1])",
+			want: 4,
+		},
+		{
+			name: "binary operations with indexing",
+			code: "IntArray[0] + IntArray[1]",
+			want: 3,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			program, err := expr.Compile(tt.code, expr.Env(mockEnv{}))
+			require.NoError(t, err, "Compile() error")
+
+			got, err := expr.Run(program, request)
+			require.NoError(t, err, "Run() error")
+
+			assert.Equal(t, tt.want, got)
+		})
+	}
+}

internal/conf/config.go

@@ -1,11 +1,46 @@
 package conf
 
-import "reflect"
+import (
+	"fmt"
+	"reflect"
+)
 
 type Config struct {
-	MapEnv bool
-	Types  TypesTable
-	Expect reflect.Kind
+	MapEnv    bool
+	Types     TypesTable
+	Operators OperatorsTable
+	Expect    reflect.Kind
+}
+
+func New(i interface{}) *Config {
+	var mapEnv bool
+	if _, ok := i.(map[string]interface{}); ok {
+		mapEnv = true
+	}
+
+	return &Config{
+		MapEnv: mapEnv,
+		Types:  CreateTypesTable(i),
+	}
+}
+
+// Check validates the compiler configuration.
+func (c *Config) Check() error {
+	// check that all functions that define operator overloading
+	// exist in environment and have correct signatures.
+	for op, fns := range c.Operators {
+		for _, fn := range fns {
+			fnType, ok := c.Types[fn]
+			if !ok || fnType.Type.Kind() != reflect.Func {
+				return fmt.Errorf("function %s for %s operator does not exist in environment", fn, op)
+			}
+
+			if fnType.Type.NumIn() != 2 || fnType.Type.NumOut() != 1 {
+				return fmt.Errorf("function %s for %s operator does not have a correct signature", fn, op)
+			}
+		}
+	}
+	return nil
 }
 
 // Option for configuring config.

internal/conf/operators_table.go

@@ -0,0 +1,5 @@
+package conf
+
+// 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

internal/conf/types_table.go

@@ -9,25 +9,17 @@ type Tag struct {
 
 type TypesTable map[string]Tag
 
-func New(i interface{}) *Config {
-	var mapEnv bool
-	if _, ok := i.(map[string]interface{}); ok {
-		mapEnv = true
-	}
-
-	return &Config{
-		MapEnv: mapEnv,
-		Types:  CreateTypesTable(i),
-	}
-}
-
 // CreateTypesTable creates types table for type checks during parsing.
 // If struct is passed, all fields will be treated as variables,
 // as well as all fields of embedded structs and struct itself.
 //
 // If map is passed, all items will be treated as variables
 // (key as name, value as type).
 func CreateTypesTable(i interface{}) TypesTable {
+	if i == nil {
+		return nil
+	}
+
 	types := make(TypesTable)
 	v := reflect.ValueOf(i)
 	t := reflect.TypeOf(i)