JITコンパイルする電卓をGoで実装してみたが正しく動いてくれない

Question

RustでJITコンパイルする電卓を実装してみる
というQiitaの記事を拝見し、同じプログラムをGo言語で実装しようとしましてひとまずセグメンテーションエラーが出ないところまでは行ったのですが、正しく動作しません。

以下のような動作になります。

$ go run *.go
>> 1
ast: 1
824634400832
>> 1 + 1
ast: (1 + 1)
824634400864
>>

JITコンパイラの箇所が間違えていることはわかるのですが、どのように修正すればよいかわかりません。

私の環境は以下になります。

MacOS: 10.15.6
Go: go version go1.15 darwin/amd64

不躾ながらソースコードをこちらに貼らせていただきます。

calculator.go : Lexer, Parser, Evalがあるコード
jit_compiler.go: JITコンパイラを実装したコード
main.go: メイン

calculator.go

package main

import (
	"bytes"
	"fmt"
	"strconv"
	"unsafe"
)

type TokenType string

const (
	ILLEGAL = "ILLEGAL"
	EOF     = "EOF"

	INT = "INT"

	// Operators
	PLUS     = "+"
	MINUS    = "-"
	ASTERISK = "*"
	SLASH    = "/"
)

type Token struct {
	Type    TokenType
	Literal string
}

type Lexer struct {
	input        string
	position     int  // current position in input (points to current char)
	readPosition int  // current reading position in input (after current char)
	ch           byte // current char under examination
}

func NewLexer(input string) *Lexer {
	l := &Lexer{input: input}
	l.readChar()
	return l
}

func (l *Lexer) NextToken() Token {
	var tok Token

	l.skipWhitespace()

	switch l.ch {
	case '+':
		tok = newToken(PLUS, l.ch)
	case '-':
		tok = newToken(MINUS, l.ch)
	case '/':
		tok = newToken(SLASH, l.ch)
	case '*':
		tok = newToken(ASTERISK, l.ch)
	case 0:
		tok.Literal = ""
		tok.Type = EOF
	default:
		if isDigit(l.ch) {
			tok.Type = INT
			tok.Literal = l.readNumber()
			return tok
		} else {
			tok = newToken(ILLEGAL, l.ch)
		}
	}

	l.readChar()
	return tok
}

func (l *Lexer) skipWhitespace() {
	for l.ch == ' ' || l.ch == '\t' || l.ch == '\n' || l.ch == '\r' {
		l.readChar()
	}
}

func (l *Lexer) readChar() {
	if l.readPosition >= len(l.input) {
		l.ch = 0
	} else {
		l.ch = l.input[l.readPosition]
	}
	l.position = l.readPosition
	l.readPosition += 1
}

func (l *Lexer) readNumber() string {
	position := l.position
	for isDigit(l.ch) {
		l.readChar()
	}
	return l.input[position:l.position]
}

func isDigit(ch byte) bool {
	return '0' <= ch && ch <= '9'
}

func newToken(tokenType TokenType, ch byte) Token {
	return Token{Type: tokenType, Literal: string(ch)}
}

// AST

type Node interface {
	TokenLiteral() string
	String() string
}

type IntegerLiteral struct {
	Token Token
	Value int64
}

func (il *IntegerLiteral) TokenLiteral() string { return il.Token.Literal }
func (il *IntegerLiteral) String() string       { return il.Token.Literal }

type InfixExpression struct {
	Token    Token // The operator token, e.g. +
	Left     Node
	Operator string
	Right    Node
}

func (oe *InfixExpression) TokenLiteral() string { return oe.Token.Literal }
func (oe *InfixExpression) String() string {
	var out bytes.Buffer

	out.WriteString("(")
	out.WriteString(oe.Left.String())
	out.WriteString(" " + oe.Operator + " ")
	out.WriteString(oe.Right.String())
	out.WriteString(")")

	return out.String()
}

// Parser

type Parser struct {
	l *Lexer

	errors []string

	curToken  Token
	peekToken Token
}

func NewParser(l *Lexer) *Parser {
	p := &Parser{
		l:      l,
		errors: []string{},
	}
	// Read Two tokens. Both curToken and peekToken are set.
	p.nextToken()
	p.nextToken()

	return p
}

func (p *Parser) nextToken() {
	p.curToken = p.peekToken
	p.peekToken = p.l.NextToken()
}

func (p *Parser) Parse() (Node, error) {
	ast, err := p.AddOrSub()
	if err != nil {
		return nil, err
	}
	return ast, nil
}

func (p *Parser) AddOrSub() (Node, error) {
	left, err := p.MulOrDiv()
	if err != nil {
		return nil, err
	}
loop:
	for p.curToken.Type != EOF {
		switch p.curToken.Type {
		case PLUS, MINUS:
			token := p.curToken
			p.nextToken()
			right, err := p.MulOrDiv()
			if err != nil {
				return nil, err
			}
			node := &InfixExpression{
				Token:    token,
				Left:     left,
				Operator: token.Literal,
				Right:    right,
			}
			left = node
		default:
			break loop
		}
	}
	return left, nil
}

func (p *Parser) MulOrDiv() (Node, error) {
	left, err := p.Num()
	if err != nil {
		return nil, err
	}
loop:
	for p.curToken.Type != EOF {
		switch p.curToken.Type {
		case ASTERISK, SLASH:
			token := p.curToken
			p.nextToken()
			right, err := p.Num()
			if err != nil {
				return nil, err
			}
			node := &InfixExpression{
				Token:    token,
				Left:     left,
				Operator: token.Literal,
				Right:    right,
			}
			left = node
		default:
			break loop
		}
	}
	return left, nil
}

func (p *Parser) Num() (Node, error) {
	if p.curToken.Type == INT {
		token := p.curToken
		p.nextToken()

		value, err := strconv.ParseInt(token.Literal, 0, 64)
		if err != nil {
			return nil, err
		}

		node := &IntegerLiteral{
			Token: token,
			Value: value,
		}

		return node, nil
	}
	return nil, fmt.Errorf("error in Num()")
}

// Eval

func Eval(ast Node) int64 {
	var useCompiler bool = true

	if useCompiler {
		compiler := NewCompiler()
		compiler.GenCode(ast)
		// fmt.Println("compiler.mmapFunc:", compiler.mmapFunc)

		type execFunc func() uint
		unsafeFunc := (uintptr)(unsafe.Pointer(&compiler.mmapFunc))
		f := *(*execFunc)(unsafe.Pointer(&unsafeFunc))
		value := f()
		return int64(value)
	} else {
		switch node := ast.(type) {
		case *IntegerLiteral:
			return node.Value
		case *InfixExpression:
			left := Eval(node.Left)
			right := Eval(node.Right)
			switch node.Operator {
			case "+":
				return left + right
			case "-":
				return left - right
			case "*":
				return left * right
			case "/":
				return left / right
			}
		}
	}
	return 0
}

jit_compiler.go

package main

import (
	"syscall"
)

type Compiler struct {
	mmapFunc   []byte
	currentIdx int
}

func NewCompiler() *Compiler {
	mmap, err := syscall.Mmap(
		-1,
		0,
		1024,
		syscall.PROT_READ|syscall.PROT_WRITE|syscall.PROT_EXEC,
		syscall.MAP_PRIVATE|syscall.MAP_ANON,
	)
	if err != nil {
		panic(err)
	}
	return &Compiler{
		mmapFunc: mmap,
	}
}

func (c *Compiler) pushCode(code []uint8) {
	for _, b := range code {
		c.mmapFunc[c.currentIdx] = b
		c.currentIdx++
	}
}

func (c *Compiler) genCodeAST(ast Node) {
	switch node := ast.(type) {
	case *IntegerLiteral:
		c.pushCode([]uint8{0x6a, uint8(node.Value)}) // push {}
	case *InfixExpression:
		c.genCodeAST(node.Left)
		c.genCodeAST(node.Right)
		c.pushCode([]uint8{0x5f}) // pop rdi
		c.pushCode([]uint8{0x58}) // pop rax
		switch node.Operator {
		case "+":
			c.pushCode([]uint8{0x48, 0x01, 0xf8}) // add rax, rdi
		case "-":
			c.pushCode([]uint8{0x48, 0x29, 0xf8}) // sub rax, rdi
		case "*":
			c.pushCode([]uint8{0x48, 0x0f, 0xaf, 0xc7}) // imul rax, rdi
		case "/":
			c.pushCode([]uint8{0x48, 0x99})       // cqo
			c.pushCode([]uint8{0x48, 0xf7, 0xff}) // idiv rdi
		}
		c.pushCode([]uint8{0x50}) // push rax
	}
}

func (c *Compiler) GenCode(ast Node) {
	c.genCodeAST(ast)
	c.pushCode([]uint8{0x58}) // pop rax
	c.pushCode([]uint8{0xc3}) // ret
}

main.go

package main

import (
	"bufio"
	"fmt"
	"io"
	"os"
)

const PROMPT = ">> "

func Start(in io.Reader, out io.Writer) {
	scanner := bufio.NewScanner(in)

	for {
		fmt.Printf(PROMPT)
		scanned := scanner.Scan()
		if !scanned {
			return
		}

		line := scanner.Text()
		l := NewLexer(line)
		p := NewParser(l)

		ast, err := p.Parse()
		if err != nil {
			panic(err)
		}
		fmt.Println("ast:", ast)

		evaluated := Eval(ast)
		fmt.Println(evaluated)
	}
}

func main() {
	Start(os.Stdin, os.Stdout)
}

参考にさせてもらっているOSSコード

• https://github.com/0yoyoyo/calculator/blob/master/src/lib.rs
• https://github.com/bspaans/jit-compiler/blob/3f22b8f79b13b3a2ccd3788e8c9f6d073812c985/lib/machinecode.go