eval expression in go

expr.go

package main

import (
	"go/ast"
	"go/parser"
	"go/token"
	"go/constant"
	"reflect"
	"errors"
	"fmt"
)


type VStack struct {
	s []constant.Value
}

type OStack struct {
	s []func()error
}

type Expr struct {
	expr ast.Expr
}

func NewExpr(src string) (*Expr, error) {
	f, err := parser.ParseExpr(src)
	if err != nil {
		return nil, err
	}
ast.Print(nil, f)
	return &Expr {
		expr: f,
	}, nil
}
var ArgsNotEnough = errors.New("args not enough")
var FuncNotExists = errors.New("func not exists")
var InvalidValue = errors.New("invalid value")

func Op2(vs *VStack, op token.Token) func() error {
	return func() error {
		v2 := vs.Pop()
		v1 := vs.Pop()
		if v1 == nil || v2 == nil {
			return ArgsNotEnough
		}
		var vr constant.Value
		switch op {
		case token.EQL, token.LSS, token.GTR, token.NEQ, token.LEQ, token.GEQ:
			vr = constant.MakeBool(constant.Compare(v1, op, v2))
		case token.SHL, token.SHR:
			sh, ok := constant.Uint64Val(v2)
			if ! ok {
				return InvalidValue
			}
			vr = constant.Shift(v1, op, uint(sh))
		default:
			vr = constant.BinaryOp(v1, op, v2)
		}
		//fmt.Printf("%v %s %v = %v\n", v1, op.String(), v2, vr)
		vs.Push(vr)
		return nil
	}
}

func Op1(vs *VStack, op token.Token) func() error {
	return func() error {
		v1 := vs.Pop()
		if v1 == nil {
			return ArgsNotEnough
		}
		vr := constant.UnaryOp(op, v1, 64)
		//fmt.Printf("%s %v = %v\n", op.String(), v1, vr)
		vs.Push(vr)
		return nil
	}
}

type UFunc func([]constant.Value)constant.Value

func FuncN(vs *VStack, n int, funcs map[string]UFunc) func() error {
	return func() error {
		args := make([]constant.Value, n)
		for i := n - 1; i >= 0; i-- {
			args[i] = vs.Pop()
			if args[i] == nil {
				return ArgsNotEnough
			}
		}
		fv := vs.Pop()
		if fv == nil {
			return FuncNotExists
		}
		fname := constant.StringVal(fv)
		f := funcs[fname]
		if f == nil {
			return FuncNotExists
		}
		vr := f(args)
		//fmt.Printf("%s(%v) = %v\n", fname, args, vr)
		vs.Push(vr)
		return nil
	}
}

func (s *VStack) Push(v constant.Value) {
	s.s = append(s.s, v)
}

func (s *VStack) Pop() constant.Value {
	l := len(s.s)
	if l <= 0 {
		return nil
	}
	ret := s.s[l - 1]
	s.s = s.s[0:l-1]
	return ret
}

func (s *OStack) Push(f func() error) {
	s.s = append(s.s, f)
}

func (s *OStack) Pop() func() error {
	l := len(s.s)
	if l <= 0 {
		return nil
	}
	ret := s.s[l - 1]
	s.s= s.s[0:l-1]
	return ret
}

func (e *Expr) Eval(vars map[string]constant.Value, funcs map[string]UFunc) (ret constant.Value, err error) {
	defer func(){
		e := recover()
		if e != nil {
			err = fmt.Errorf("ERR: %v", e)
		}
	}()

	vs := VStack{}
	os := OStack{}
	isV := 0
	ast.Inspect(e.expr, func (n ast.Node) (ret bool) {
		switch x := n.(type) {
		case nil:
			if isV > 0 {
				isV--
				return true
			} else {
				err = os.Pop()()
				if err != nil {
					return false
				}
			}
		case *ast.BasicLit:
			v := constant.MakeFromLiteral(x.Value, x.Kind, 0)
			vs.Push(v)
			isV++
		case *ast.Ident:
			v := vars[x.Name]
			if v == nil {
				v = constant.MakeString(x.Name)
			}
			vs.Push(v)
			isV++
		case *ast.BinaryExpr:
			os.Push(Op2(&vs, x.Op))
		case *ast.UnaryExpr:
			os.Push(Op1(&vs, x.Op))
		case *ast.ParenExpr:
			isV++
			//s = "()"
		case *ast.CallExpr:
			n := len(x.Args)
			os.Push(FuncN(&vs, n, funcs))
	/*
			case *ast.SelectorExpr:
				s = "."
			case *ast.IndexExpr:
				s = "[]"
			case *ast.CompositeLit:
				s = "{}"
			case *ast.MapType:
				s = "map"
			case *ast.ArrayType:
				s = "arr"
			case *ast.KeyValueExpr:
				s = "=>"
			case *ast.StarExpr:
				s = "*->"
			case *ast.StructType:
				s = "struct"
			case *ast.FieldList:
				s = ":"
			case *ast.Field:
				s = "FIELD"
			case *ast.FuncLit:
				s = "func"
	*/
		default:
			s := reflect.TypeOf(x).String()
			panic("bad token: "+s)
		}
		return true
	})
	ret = vs.Pop()
	return
}

func main() {
	var vars = make(map[string]constant.Value)
	var funcs = make(map[string]UFunc)
	vars["id"] = constant.MakeInt64(10086)
	vars["true"] = constant.MakeBool(true)
	vars["false"] = constant.MakeBool(false)
	funcs["add"] = func(args []constant.Value)constant.Value {
		v := constant.MakeInt64(0)
		for _, a := range args {
			v = constant.BinaryOp(v, token.ADD, a)
		}
		return v
	}

	src := `(id % 1024) + (func(a,b int){ return a+b})(1,2)`
	fmt.Println(src)
	expr, err := NewExpr(src)
	if err != nil {
		panic(err)
	}
	ret, err := expr.Eval(vars, funcs)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(ret)
}