Add expr.Function feature

Author: antonmedv

ast/node.go

@@ -51,8 +51,8 @@ type IdentifierNode struct {
 	Value       string
 	Deref       bool
 	FieldIndex  []int
-	Method      bool
-	MethodIndex int
+	Method      bool // true if method, false if field
+	MethodIndex int  // index of method, set only if Method is true
 }
 
 type IntegerNode struct {
@@ -124,6 +124,7 @@ type CallNode struct {
 	Arguments []Node
 	Typed     int
 	Fast      bool
+	Func      func(params ...interface{}) (interface{}, error)
 }
 
 type BuiltinNode struct {

checker/checker.go

@@ -60,6 +60,7 @@ type visitor struct {
 
 type info struct {
 	method bool
+	fn     *conf.Function
 }
 
 func (v *visitor) visit(node ast.Node) (reflect.Type, info) {
@@ -134,6 +135,12 @@ func (v *visitor) IdentifierNode(node *ast.IdentifierNode) (reflect.Type, info)
 		node.Deref = true
 		return anyType, info{}
 	}
+	if fn, ok := v.config.Functions[node.Value]; ok {
+		// Return anyType instead of func type as we don't know the arguments yet.
+		// The func type can be one of the fn.Types. The type will be resolved
+		// when the arguments are known in CallNode.
+		return anyType, info{fn: fn}
+	}
 	if t, ok := v.config.Types[node.Value]; ok {
 		if t.Ambiguous {
 			return v.error(node, "ambiguous identifier %v", node.Value)
@@ -466,6 +473,32 @@ func (v *visitor) SliceNode(node *ast.SliceNode) (reflect.Type, info) {
 func (v *visitor) CallNode(node *ast.CallNode) (reflect.Type, info) {
 	fn, fnInfo := v.visit(node.Callee)
 
+	if fnInfo.fn != nil {
+		f := fnInfo.fn
+		node.Func = f.Func
+		if len(f.Types) == 0 {
+			// No type was specified, so we assume the function returns any.
+			return anyType, info{}
+		}
+		var firstErr *file.Error
+		for _, t := range f.Types {
+			outType, err := v.checkFunc(f.Name, t, false, node)
+			if err != nil {
+				if firstErr == nil {
+					firstErr = err
+				}
+				continue
+			}
+			return outType, info{}
+		}
+		if firstErr != nil {
+			if v.err == nil {
+				v.err = firstErr
+			}
+			return anyType, info{}
+		}
+	}
+
 	fnName := "function"
 	if identifier, ok := node.Callee.(*ast.IdentifierNode); ok {
 		fnName = identifier.Value
@@ -475,7 +508,6 @@ func (v *visitor) CallNode(node *ast.CallNode) (reflect.Type, info) {
 			fnName = name.Value
 		}
 	}
-
 	switch fn.Kind() {
 	case reflect.Interface:
 		return anyType, info{}
@@ -484,37 +516,50 @@ func (v *visitor) CallNode(node *ast.CallNode) (reflect.Type, info) {
 		if fnInfo.method {
 			inputParamsCount = 2 // for methods
 		}
-
+		// TODO: Deprecate OpCallFast and move fn(...any) any to TypedFunc list.
+		// To do this we need add support for variadic arguments in OpCallTyped.
 		if !isAny(fn) &&
 			fn.IsVariadic() &&
 			fn.NumIn() == inputParamsCount &&
-			((fn.NumOut() == 1 && // Function with one return value
-				fn.Out(0).Kind() == reflect.Interface) ||
-				(fn.NumOut() == 2 && // Function with one return value and an error
-					fn.Out(0).Kind() == reflect.Interface &&
-					fn.Out(1) == errorType)) {
+			fn.NumOut() == 1 &&
+			fn.Out(0).Kind() == reflect.Interface {
 			rest := fn.In(fn.NumIn() - 1) // function has only one param for functions and two for methods
 			if rest.Kind() == reflect.Slice && rest.Elem().Kind() == reflect.Interface {
 				node.Fast = true
 			}
 		}
 
-		return v.checkFunc(fn, fnInfo.method, node, fnName, node.Arguments)
+		outType, err := v.checkFunc(fnName, fn, fnInfo.method, node)
+		if err != nil {
+			if v.err == nil {
+				v.err = err
+			}
+			return anyType, info{}
+		}
+
+		v.findTypedFunc(node, fn, fnInfo.method)
+
+		return outType, info{}
 	}
 	return v.error(node, "%v is not callable", fn)
 }
 
-// checkFunc checks func arguments and returns "return type" of func or method.
-func (v *visitor) checkFunc(fn reflect.Type, method bool, node *ast.CallNode, name string, arguments []ast.Node) (reflect.Type, info) {
+func (v *visitor) checkFunc(name string, fn reflect.Type, method bool, node *ast.CallNode) (reflect.Type, *file.Error) {
 	if isAny(fn) {
-		return anyType, info{}
+		return anyType, nil
 	}
 
 	if fn.NumOut() == 0 {
-		return v.error(node, "func %v doesn't return value", name)
+		return anyType, &file.Error{
+			Location: node.Location(),
+			Message:  fmt.Sprintf("func %v doesn't return value", name),
+		}
 	}
 	if numOut := fn.NumOut(); numOut > 2 {
-		return v.error(node, "func %v returns more then two values", name)
+		return anyType, &file.Error{
+			Location: node.Location(),
+			Message:  fmt.Sprintf("func %v returns more then two values", name),
+		}
 	}
 
 	// If func is method on an env, first argument should be a receiver,
@@ -530,19 +575,28 @@ func (v *visitor) checkFunc(fn reflect.Type, method bool, node *ast.CallNode, na
 	}
 
 	if fn.IsVariadic() {
-		if len(arguments) < fnNumIn-1 {
-			return v.error(node, "not enough arguments to call %v", name)
+		if len(node.Arguments) < fnNumIn-1 {
+			return anyType, &file.Error{
+				Location: node.Location(),
+				Message:  fmt.Sprintf("not enough arguments to call %v", name),
+			}
 		}
 	} else {
-		if len(arguments) > fnNumIn {
-			return v.error(node, "too many arguments to call %v", name)
+		if len(node.Arguments) > fnNumIn {
+			return anyType, &file.Error{
+				Location: node.Location(),
+				Message:  fmt.Sprintf("too many arguments to call %v", name),
+			}
 		}
-		if len(arguments) < fnNumIn {
-			return v.error(node, "not enough arguments to call %v", name)
+		if len(node.Arguments) < fnNumIn {
+			return anyType, &file.Error{
+				Location: node.Location(),
+				Message:  fmt.Sprintf("not enough arguments to call %v", name),
+			}
 		}
 	}
 
-	for i, arg := range arguments {
+	for i, arg := range node.Arguments {
 		t, _ := v.visit(arg)
 
 		var in reflect.Type
@@ -564,44 +618,14 @@ func (v *visitor) checkFunc(fn reflect.Type, method bool, node *ast.CallNode, na
 		}
 
 		if !t.AssignableTo(in) && t.Kind() != reflect.Interface {
-			return v.error(arg, "cannot use %v as argument (type %v) to call %v ", t, in, name)
-		}
-	}
-
-	// OnCallTyped doesn't work for functions with variadic arguments,
-	// and doesn't work named function, like `type MyFunc func() int`.
-	// In PkgPath() is an empty string, it's unnamed function.
-	if !fn.IsVariadic() && fn.PkgPath() == "" {
-	funcTypes:
-		for i := range vm.FuncTypes {
-			if i == 0 {
-				continue
-			}
-			typed := reflect.ValueOf(vm.FuncTypes[i]).Elem().Type()
-			if typed.Kind() != reflect.Func {
-				continue
-			}
-			if typed.NumOut() != fn.NumOut() {
-				continue
+			return anyType, &file.Error{
+				Location: arg.Location(),
+				Message:  fmt.Sprintf("cannot use %v as argument (type %v) to call %v ", t, in, name),
 			}
-			for j := 0; j < typed.NumOut(); j++ {
-				if typed.Out(j) != fn.Out(j) {
-					continue funcTypes
-				}
-			}
-			if typed.NumIn() != fnNumIn {
-				continue
-			}
-			for j := 0; j < typed.NumIn(); j++ {
-				if typed.In(j) != fn.In(j+fnInOffset) {
-					continue funcTypes
-				}
-			}
-			node.Typed = i
 		}
 	}
 
-	return fn.Out(0), info{}
+	return fn.Out(0), nil
 }
 
 func (v *visitor) BuiltinNode(node *ast.BuiltinNode) (reflect.Type, info) {
@@ -769,3 +793,44 @@ func (v *visitor) PairNode(node *ast.PairNode) (reflect.Type, info) {
 	v.visit(node.Value)
 	return nilType, info{}
 }
+
+func (v *visitor) findTypedFunc(node *ast.CallNode, fn reflect.Type, method bool) {
+	// OnCallTyped doesn't work for functions with variadic arguments,
+	// and doesn't work named function, like `type MyFunc func() int`.
+	// In PkgPath() is an empty string, it's unnamed function.
+	if !fn.IsVariadic() && fn.PkgPath() == "" {
+		fnNumIn := fn.NumIn()
+		fnInOffset := 0
+		if method {
+			fnNumIn--
+			fnInOffset = 1
+		}
+	funcTypes:
+		for i := range vm.FuncTypes {
+			if i == 0 {
+				continue
+			}
+			typed := reflect.ValueOf(vm.FuncTypes[i]).Elem().Type()
+			if typed.Kind() != reflect.Func {
+				continue
+			}
+			if typed.NumOut() != fn.NumOut() {
+				continue
+			}
+			for j := 0; j < typed.NumOut(); j++ {
+				if typed.Out(j) != fn.Out(j) {
+					continue funcTypes
+				}
+			}
+			if typed.NumIn() != fnNumIn {
+				continue
+			}
+			for j := 0; j < typed.NumIn(); j++ {
+				if typed.In(j) != fn.In(j+fnInOffset) {
+					continue funcTypes
+				}
+			}
+			node.Typed = i
+		}
+	}
+}

compiler/compiler.go

@@ -2,7 +2,6 @@ package compiler
 
 import (
 	"fmt"
-	"math"
 	"reflect"
 
 	"github.com/antonmedv/expr/ast"
@@ -25,8 +24,9 @@ func Compile(tree *parser.Tree, config *conf.Config) (program *Program, err erro
 	}()
 
 	c := &compiler{
-		index:     make(map[interface{}]int),
-		locations: make([]file.Location, 0),
+		locations:      make([]file.Location, 0),
+		constantsIndex: make(map[interface{}]int),
+		functionsIndex: make(map[string]int),
 	}
 
 	if config != nil {
@@ -52,20 +52,23 @@ func Compile(tree *parser.Tree, config *conf.Config) (program *Program, err erro
 		Constants: c.constants,
 		Bytecode:  c.bytecode,
 		Arguments: c.arguments,
+		Functions: c.functions,
 	}
 	return
 }
 
 type compiler struct {
-	locations []file.Location
-	constants []interface{}
-	bytecode  []Opcode
-	index     map[interface{}]int
-	mapEnv    bool
-	cast      reflect.Kind
-	nodes     []ast.Node
-	chains    [][]int
-	arguments []int
+	locations      []file.Location
+	bytecode       []Opcode
+	constants      []interface{}
+	constantsIndex map[interface{}]int
+	functions      []Function
+	functionsIndex map[string]int
+	mapEnv         bool
+	cast           reflect.Kind
+	nodes          []ast.Node
+	chains         [][]int
+	arguments      []int
 }
 
 func (c *compiler) emitLocation(loc file.Location, op Opcode, arg int) int {
@@ -110,25 +113,30 @@ func (c *compiler) addConstant(constant interface{}) int {
 		indexable = true
 		hash = fmt.Sprintf("%v", method)
 	}
-
 	if indexable {
-		if p, ok := c.index[hash]; ok {
+		if p, ok := c.constantsIndex[hash]; ok {
 			return p
 		}
 	}
-
 	c.constants = append(c.constants, constant)
-	if len(c.constants) > math.MaxUint16 {
-		panic("exceeded constants max space limit")
-	}
-
 	p := len(c.constants) - 1
 	if indexable {
-		c.index[hash] = p
+		c.constantsIndex[hash] = p
 	}
 	return p
 }
 
+func (c *compiler) addFunction(f Function) int {
+	addr := fmt.Sprintf("%v", f)
+	if p, ok := c.functionsIndex[addr]; ok {
+		return p
+	}
+	p := len(c.functions)
+	c.functions = append(c.functions, f)
+	c.functionsIndex[addr] = p
+	return p
+}
+
 func (c *compiler) patchJump(placeholder int) {
 	offset := len(c.bytecode) - placeholder
 	c.arguments[placeholder-1] = offset
@@ -516,6 +524,22 @@ func (c *compiler) CallNode(node *ast.CallNode) {
 	for _, arg := range node.Arguments {
 		c.compile(arg)
 	}
+	if node.Func != nil {
+		switch len(node.Arguments) {
+		case 0:
+			c.emit(OpCall0, c.addFunction(node.Func))
+		case 1:
+			c.emit(OpCall1, c.addFunction(node.Func))
+		case 2:
+			c.emit(OpCall2, c.addFunction(node.Func))
+		case 3:
+			c.emit(OpCall3, c.addFunction(node.Func))
+		default:
+			c.emit(OpLoadFunc, c.addFunction(node.Func))
+			c.emit(OpCallN, len(node.Arguments))
+		}
+		return
+	}
 	c.compile(node.Callee)
 	if node.Typed > 0 {
 		c.emit(OpCallTyped, node.Typed)