neo-go/pkg/vm/compiler/compiler.go

package compiler

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"go/ast"
	"go/build"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"text/tabwriter"

	"github.com/CityOfZion/neo-go/pkg/vm"
)

const fileExt = "avm"

// Options contains all the parameters that affect the behaviour of the compiler.
type Options struct {
	// The extension of the output file default set to .avm
	Ext string

	// The name of the output file.
	Outfile string

	// Debug will output an hex encoded string of the generated bytecode.
	Debug bool
}

// Compile compiles a Go program into bytecode that can run on the NEO virtual machine.
func Compile(input io.Reader, o *Options) ([]byte, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "", input, 0)
	if err != nil {
		return nil, err
	}

	conf := types.Config{Importer: importer.Default()}
	typeInfo := &types.Info{
		Types:      make(map[ast.Expr]types.TypeAndValue),
		Defs:       make(map[*ast.Ident]types.Object),
		Uses:       make(map[*ast.Ident]types.Object),
		Implicits:  make(map[ast.Node]types.Object),
		Selections: make(map[*ast.SelectorExpr]*types.Selection),
		Scopes:     make(map[ast.Node]*types.Scope),
	}

	// Typechecker
	_, err = conf.Check("", fset, []*ast.File{f}, typeInfo)
	if err != nil {
		return nil, err
	}

	imports, err := resolveImports(f)
	if err != nil {
		return nil, err
	}

	buf, err := CodeGen(f, typeInfo, imports)
	if err != nil {
		return nil, err
	}

	return buf.Bytes(), nil
}

type archive struct {
	f        *ast.File
	typeInfo *types.Info
}

func resolveImports(f *ast.File) (map[string]*archive, error) {
	packages := map[string]*archive{}
	for _, imp := range f.Imports {
		path := strings.Replace(imp.Path.Value, `"`, "", 2)
		path = filepath.Join(gopath(), "src", path)
		fset := token.NewFileSet()
		pkgs, err := parser.ParseDir(fset, path, nil, 0)
		if err != nil {
			return nil, err
		}

		for name, pkg := range pkgs {
			file := ast.MergePackageFiles(pkg, 0)
			conf := types.Config{Importer: importer.Default()}
			typeInfo := &types.Info{
				Types:      make(map[ast.Expr]types.TypeAndValue),
				Defs:       make(map[*ast.Ident]types.Object),
				Uses:       make(map[*ast.Ident]types.Object),
				Implicits:  make(map[ast.Node]types.Object),
				Selections: make(map[*ast.SelectorExpr]*types.Selection),
				Scopes:     make(map[ast.Node]*types.Scope),
			}

			// Typechecker
			_, err = conf.Check("", fset, []*ast.File{file}, typeInfo)
			if err != nil {
				return nil, err
			}
			packages[name] = &archive{file, typeInfo}
		}
	}
	return packages, nil
}

// CompileAndSave will compile and save the file to disk.
func CompileAndSave(src string, o *Options) error {
	if len(o.Outfile) == 0 {
		if !strings.HasSuffix(src, ".go") {
			return errors.New("not a Go file")
		}
		o.Outfile = strings.TrimSuffix(src, ".go")
	}
	if len(o.Ext) == 0 {
		o.Ext = fileExt
	}
	b, err := ioutil.ReadFile(src)
	if err != nil {
		return err
	}
	b, err = Compile(bytes.NewReader(b), o)
	if err != nil {
		return err
	}
	if o.Debug {
		log.Println(hex.EncodeToString(b))
	}

	out := fmt.Sprintf("%s.%s", o.Outfile, o.Ext)
	return ioutil.WriteFile(out, b, os.ModePerm)
}

// DumpOpcode compiles the program and dumps the opcode in a user friendly format.
func DumpOpcode(src string) error {
	b, err := ioutil.ReadFile(src)
	if err != nil {
		return err
	}
	b, err = Compile(bytes.NewReader(b), &Options{})
	if err != nil {
		return err
	}

	w := tabwriter.NewWriter(os.Stdout, 0, 0, 4, ' ', 0)
	fmt.Fprintln(w, "INDEX\tOPCODE\tDESC\t")
	for i := 0; i < len(b); i++ {
		fmt.Fprintf(w, "%d\t0x%2x\t%s\t\n", i, b[i], vm.Opcode(b[i]))
	}
	w.Flush()
	return nil
}

func gopath() string {
	gopath := os.Getenv("GOPATH")
	if len(gopath) == 0 {
		gopath = build.Default.GOPATH
	}
	return gopath
}

func init() {
	log.SetFlags(0)
}