簡単な式の評価機を作ってみる #golang

はじめに

皆さんはgoパッケージを使ったことはありますか？私は、Go！Go!言いながら、恥ずかしながら、あまりまともに使ったことがありませんでした。最近は、仕事でもAST（抽象構文木）を弄り倒すことがあるので、今回はgoパッケージに触れたいと思います。

ここでは、簡単な四則演算などを評価し、結果を返すREPLを作ってみます。一見難しそうですが、ほとんどgoパッケージの機能を使うので、非常に簡単です。

ASTへパースする

まずは入力された文字列をパースし、ASTを生成してみます。ここでは、式の評価なので式単位でパースします。パースには、parser.ParseExprを用います。
たとえば、以下のように使用できます。

expr, err := parser.ParseExpr("1+1")
if err != nil {
	panic(err)
}

ast.Inspect(expr, func(n ast.Node) bool {
	if n != nil {
		fmt.Printf("%[1]v(%[1]T)\n", n)
	}
	return true
})

Playgroundで動かす

なお、ast.InspectはASTのノードをトラバースする関数です。この場合、1+1をパースすると、以下のようなASTが取得できます。

*ast.BinaryExpr (+)
├── *ast.BasicLit (1)
└── *ast.BasicLit (1)

このASTを評価することで、式の評価を実現します。さて、具体的に評価とはどうすれば良いのでしょうか？

1+1の評価

上記の例でみたとおり、2項演算を行う式は、ast.BinaryExprで表されます。まずは2項演算式である1+1を評価してみましょう。

func eval1pls1(expr *ast.BinaryExpr) (int64, error) {
	xLit, ok := expr.X.(*ast.BasicLit)
	if !ok {
		return 0, errors.New("left operand is not BasicLit")
	}

	yLit, ok := expr.Y.(*ast.BasicLit)
	if !ok {
		return 0, errors.New("right operand is not BasicLit")
	}

	if expr.Op != token.ADD {
		return 0, errors.New("operator is not +")
	}

	x, err := strconv.ParseInt(xLit.Value, 10, 64)
	if err != nil {
		return 0, err
	}

	y, err := strconv.ParseInt(yLit.Value, 10, 64)
	if err != nil {
		return 0, err
	}

	return x + y, nil
}

func main() {
	expr, err := parser.ParseExpr("1+1")
	if err != nil {
		panic(err)
	}

	if be, ok := expr.(*ast.BinaryExpr); ok {
		if v, err := eval1pls1(be); err == nil {
			fmt.Println(v)
		} else {
			fmt.Println("Error:", err)
		}
	}

}

Playgroundで動かす

parser.Parseが返すast.Exprは式であることを表すインタフェースであるため、具体的に「何の」式なのかという値は型アサーション（まはた型スイッチ）しないと手にはいりません。そのため、今回必要なast.BinaryExprとその右辺と左辺に設定されているast.BasicLitを取得するために、各々型アサーションを行っています。

ast.BasicLitは基本型のリテラルを表す型で、1や"hoge"などはこの型で表されます。
一方で、trueとfalseについては、識別子であるためast.BasicLitではなく、ast.Identとして表されるので注意してください。ちなみに、trueとfalseは識別子として使用できます。
BasicLit.Valueはそのリテラルを文字列で表したもので、 42, 0x7f, 3.14, 1e-9, 2.4i, 'a', "foo" などです。なお、文字列には"などで括られているという点に注意してください。これらを除去するには、strconv.Unquoteが使えます。

2項演算式とリテラルの評価

さて、1+1は評価できるようになったものの、ここで作りたいのは2項演算式の評価です。そのため、もう少し汎用化が必要そうです。オペランドが数値であるという点とオペレータが+であるという点をどうにか汎用化しないといけません。

そこで、2項演算子を簡単に評価するために、go/constantパッケージにあるBinaryOp関数を使いましょう。

BinaryOp関数は以下のようなシグニチャを持っています。

func BinaryOp(x Value, op token.Token, y Value) Value

opはtoken.ADDなどの演算子のトークンを表すもので、これはast.BinaryExprのOpフィールドから取得できそうです。xとyに関しては、Valueフィールドからconstant.MakeFromLiteralを使って取得できます。

constant.MakeFromLiteralは以下のような定義になっています。

func MakeFromLiteral(lit string, tok token.Token, zero uint) Value

第1引数はast.BasicLitのValueフィールドを第2引数をast.BasicLitのKindフィールドを渡せば良さそうです。なお、第3引数は常に0を渡す必要があるようです。

まとめると、上記のコードは以下のように変更することができます。

func evalBinaryExpr(expr *ast.BinaryExpr) (constant.Value, error) {

	xLit, ok := expr.X.(*ast.BasicLit)
	if !ok {
		return constant.MakeUnknown(), errors.New("left operand is not BasicLit")
	}

	yLit, ok := expr.Y.(*ast.BasicLit)
	if !ok {
		return constant.MakeUnknown(), errors.New("right operand is not BasicLit")
	}

	x := evalBasicLit(xLit)
	y := evalBasicLit(yLit)
	return constant.BinaryOp(x, expr.Op, y), nil
}

func evalBasicLit(expr *ast.BasicLit) constant.Value {
	return constant.MakeFromLiteral(expr.Value, expr.Kind, 0)
}

func main() {
	expr, err := parser.ParseExpr("1+1")
	if err != nil {
		panic(err)
	}

	if be, ok := expr.(*ast.BinaryExpr); ok {
		if v, err := evalBinaryExpr(be); err == nil {
			fmt.Println(v)
		} else {
			fmt.Println("Error:", err)
		}
	}

}

Playgroundで動かす

なお、2項演算式を評価する関数をevalBinaryExprとし、リテラルの評価をする関数をevalBasicLitとしました。1-1が実行できるかチェックしてみましょう。また、"hoge"+"hoge"などもちゃんと動くことが分かると思います。

constant.MakeUnknown関数は、エラー時に使われる値を作る関数です。演算が失敗した場合にこの値になるため、ここではエラーが発生した際に返しています。なお、戻り値をconstant.Valueにしている点が重要で、後述のeval関数を作る際に効果が出ます。

式の評価

2項演算式を評価する関数を作りました。しかし、本当にやりたいのは式の評価です。式には-1などの単項演算式もありますし、1などのリテラル、(1+1)*2などの括弧付きの式もあります。もちろん、Goの文法には関数呼び出しとかもありますが、ここでは四則演算などができれば十分でしょう。

どういう種類の式なのかは、ast.Exprを具体的な型に変換することで取得することができ、その種類によって処理を分けてやれば良さそうです。

上記の2項演算式だけをやる関数を少しだけ汎用化してみましょう。
まずは式（ast.Expr）を評価する関数を作りましょう。

func evalExpr(expr ast.Expr) (v constant.Value, rerr error) {
	defer func() {
		if r := recover(); r != nil {
			v, rerr = constant.MakeUnknown(), fmt.Errorf("%v", r)
		}
	}()

	switch e := expr.(type) {
	case *ast.BinaryExpr:
		return evalBinaryExpr(e)
	case *ast.BasicLit:
		return constant.MakeFromLiteral(e.Value, e.Kind, 0), nil
	}

	return constant.MakeUnknown(), errors.New("unkown node")
}

evalExpr関数では、ast.Exprを受取り、型によって処理を分岐しています。
ここでは、2項演算式とリテラルだけ処理ができるようになっています。

deferでrecoverしているのは、constantパッケージのBinaryOpなどが型が合わなかった場合などに、エラーを返さず、panicを起こすためです。

せっかくなので、evalBinaryExpr関数で呼び出していたevalBasicLitもevalExprにおまかせしましょう。

func evalBinaryExpr(expr *ast.BinaryExpr) (constant.Value, error) {
	x, err := evalExpr(expr.X)
	if err != nil {
		return constant.MakeUnknown(), err
	}

	y, err := evalExpr(expr.Y)
	if err != nil {
		return constant.MakeUnknown(), err
	}

	return constant.BinaryOp(x, expr.Op, y), nil
}

これでだいぶスッキリしてきました。あとはevalExprに処理できる型をどんどん追加していけば良さそうです。必要そうなものを追加していくと、evalExprは以下のようになりました。

func evalExpr(expr ast.Expr) (v constant.Value, rerr error) {
	defer func() {
		if r := recover(); r != nil {
			v, rerr = constant.MakeUnknown(), fmt.Errorf("%v", r)
		}
	}()

	switch e := expr.(type) {
	case *ast.ParenExpr:
		return evalExpr(e.X)
	case *ast.BinaryExpr:
		return evalBinaryExpr(e)
	case *ast.UnaryExpr:
		return evalUnaryExpr(e)
	case *ast.BasicLit:
		return constant.MakeFromLiteral(e.Value, e.Kind, 0), nil
	case *ast.Ident:
		return evalIdent(e)
	}

	return constant.MakeUnknown(), errors.New("unkown node")
}

ast.ParenExprは、()で括られた式で、単純に中の式(ParenExpr.X)を評価してやれば良いでしょう。

ast.UnaryExprは、単項演算子式を表す型で、ast.BinaryExprと同様にevalUnaryExpr関数を作り、処理してやります。

func evalUnaryExpr(expr *ast.UnaryExpr) (constant.Value, error) {
	x, err := evalExpr(expr.X)
	if err != nil {
		return constant.MakeUnknown(), err
	}

	return constant.UnaryOp(expr.Op, x, 0), nil
}

四則演算には関係ないですが、せっかくなのでtrueやfalseも扱えるようにしましょう。trueやfalseが来た場合に、ast.Identとして処理する必要があるため、evalIdentを作ります。

func evalIdent(expr *ast.Ident) (constant.Value, error) {
	switch {
	case expr.Name == "true":
		return constant.MakeBool(true), nil
	case expr.Name == "false":
		return constant.MakeBool(false), nil
	}

	return constant.MakeUnknown(), errors.New("unkown ident")
}

trueやfalseは名前がtrueとfalseの識別子として認識されるため、上記のように処理してやれば十分です。

さて、これで完成で良さそうですが、trueとfalseを扱えるようにしたことで一つ問題があります。constant.BinaryOpでは&&や>=などの関係演算子は扱えません。代わりにconstant.Compareを使ってやる必要があります。そのため、evalBinaryExpr関数は以下のように修正する必要があります。

func evalBinaryExpr(expr *ast.BinaryExpr) (constant.Value, error) {
	x, err := evalExpr(expr.X)
	if err != nil {
		return constant.MakeUnknown(), err
	}

	y, err := evalExpr(expr.Y)
	if err != nil {
		return constant.MakeUnknown(), err
	}

	switch expr.Op {
	case token.EQL, token.NEQ, token.LSS, token.LEQ, token.GTR, token.GEQ:
		return constant.MakeBool(constant.Compare(x, expr.Op, y)), nil
	}

	return constant.BinaryOp(x, expr.Op, y), nil
}

これで、式を評価できるevalExprができました。最後にこれをラップしてeval関数を作り扱いやすくしましょう。

func eval(e string) (string, error) {
	expr, err := parser.ParseExpr(e)
	if err != nil {
		return "", err
	}

	v, err := evalExpr(expr)
	if err != nil {
		return "", err
	}

	return v.String(), nil
}

REPLを作る

式を評価する関数を作れたので、次はREPLを作ります。ここで作るREPLは式の入力を受付け、式が入力されたらそれを評価して表示するというものにします。

標準入力(os.Stdin)からbufio.Scannerを使って1行ずつ読み取るコードを書きましょう。ここではbyeと入力すればREPLを抜けれるようにします。

func repl(r io.Reader) {
	s := bufio.NewScanner(r)
	for {
		fmt.Print(">")
		if !s.Scan() {
			break
		}

		l := s.Text()
		switch {
		case l == "bye":
			return
		default:
			r, err := eval(l)
			if err != nil {
				fmt.Println("Error:", err)
				continue
			}
			fmt.Println(r)
		}
	}

	if err := s.Err(); err != nil {
		fmt.Println("Error:", err)
	}
}

うまくできましたね。ここまで作ったコードをまとめて載せておきます。

package main

import (
	"bufio"
	"errors"
	"fmt"
	"go/ast"
	"go/constant"
	"go/parser"
	"go/token"
	"io"
	"os"
)

func eval(e string) (string, error) {
	expr, err := parser.ParseExpr(e)
	if err != nil {
		return "", err
	}

	v, err := evalExpr(expr)
	if err != nil {
		return "", err
	}

	return v.String(), nil
}

func evalExpr(expr ast.Expr) (v constant.Value, rerr error) {
	defer func() {
		if r := recover(); r != nil {
			v, rerr = constant.MakeUnknown(), fmt.Errorf("%v", r)
		}
	}()

	switch e := expr.(type) {
	case *ast.ParenExpr:
		return evalExpr(e.X)
	case *ast.BinaryExpr:
		return evalBinaryExpr(e)
	case *ast.UnaryExpr:
		return evalUnaryExpr(e)
	case *ast.BasicLit:
		return constant.MakeFromLiteral(e.Value, e.Kind, 0), nil
	case *ast.Ident:
		return evalIdent(e)
	}

	return constant.MakeUnknown(), errors.New("unkown node")
}

func evalBinaryExpr(expr *ast.BinaryExpr) (constant.Value, error) {
	x, err := evalExpr(expr.X)
	if err != nil {
		return constant.MakeUnknown(), err
	}

	y, err := evalExpr(expr.Y)
	if err != nil {
		return constant.MakeUnknown(), err
	}

	switch expr.Op {
	case token.EQL, token.NEQ, token.LSS, token.LEQ, token.GTR, token.GEQ:
		return constant.MakeBool(constant.Compare(x, expr.Op, y)), nil
	}

	return constant.BinaryOp(x, expr.Op, y), nil
}

func evalUnaryExpr(expr *ast.UnaryExpr) (constant.Value, error) {
	x, err := evalExpr(expr.X)
	if err != nil {
		return constant.MakeUnknown(), err
	}

	return constant.UnaryOp(expr.Op, x, 0), nil
}

func evalIdent(expr *ast.Ident) (constant.Value, error) {
	switch {
	case expr.Name == "true":
		return constant.MakeBool(true), nil
	case expr.Name == "false":
		return constant.MakeBool(false), nil
	}

	return constant.MakeUnknown(), errors.New("unkown ident")
}

func repl(r io.Reader) {
	s := bufio.NewScanner(r)
	for {
		fmt.Print(">")
		if !s.Scan() {
			break
		}

		l := s.Text()
		switch {
		case l == "bye":
			return
		default:
			r, err := eval(l)
			if err != nil {
				fmt.Println("Error:", err)
				continue
			}
			fmt.Println(r)
		}
	}

	if err := s.Err(); err != nil {
		fmt.Println("Error:", err)
	}
}

func main() {
	repl(os.Stdin)
}

まとめ

120行で足らずで結構高性能な式の評価を行うREPLができました。
ちなみに、Goでは複素数が使えるので、今回作ったREPLは複素数も扱えます。

eval.Exprに変数(ast.Ident）や関数（ast.CallExpr）を評価する処理を追加してやれば、もっと高性能なものが作れます。
それにしても、goパッケージを使うと案外簡単に作れますね。すばらしいですね。

さて、これが実際の業務で使えるんだろうかと思ってる方もいるかもしれませんが、それは使いようだと思いますので、ぜひ用途を考えてみて下さい。