main.go

package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"io"
	"io/ioutil"
	"os"
	"path"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"text/template"

	"golang.org/x/tools/go/ast/astutil"

	"gitlab.lrde.epita.fr/spot/go2pins/channel"
	"gitlab.lrde.epita.fr/spot/go2pins/decl"
	"gitlab.lrde.epita.fr/spot/go2pins/goroutine"
	"gitlab.lrde.epita.fr/spot/go2pins/tools"
	"gitlab.lrde.epita.fr/spot/go2pins/transform"
	"gitlab.lrde.epita.fr/spot/go2pins/transform/cfg"
)

// Generate a list of cgo CStrings from the provided string slice.
func variableNamesToC(variableNames []string) []ast.Expr {
	vn := make([]ast.Expr, len(variableNames))
	for i, v := range variableNames {
		vn[i] = &ast.CallExpr{
			Fun: &ast.SelectorExpr{
				X:   ast.NewIdent("C"),
				Sel: ast.NewIdent("CString"),
			},
			Args: []ast.Expr{
				&ast.BasicLit{
					Kind:  token.STRING,
					Value: strconv.Quote(v),
				},
			},
		}
	}
	return vn
}

// Compiles the input file `(path, src)` and writes it to `file`.
func compileTo(path string, src []byte, file *os.File,
	goroutines []*goroutine.Goroutine, rers *string) (string, error) {
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, path, src, 0)

	if err != nil {
		return "", err
	}

	var config types.Config
	// FIXME: figure out if this is really needed
	config.Importer = importer.Default()
	info := &types.Info{
		Uses:       make(map[*ast.Ident]types.Object),
		Selections: make(map[*ast.SelectorExpr]*types.Selection),
	}

	// Ensure that our program type checks.
	_, err = config.Check(path, fset, []*ast.File{f}, info)

	// Apply all transformations.
	meta := transform.NewMeta(fset, info)
	res := transform.Apply(f, meta, &transform.TypeChecker{Blackbox: (*ignorefunc != "")})

	if *ignorefunc != "" || len(goroutines) != 0 {
		newsrc, blackbox := tools.Isolate(flag.Args()[0], *output, toIgnore)
		src = []byte(newsrc)
		f, err = parser.ParseFile(fset, path, src, 0)
		if err != nil {
			return "", err
		}
		if blackbox != "" {
			if err = os.MkdirAll(*output+"/blackbox", 0777); err != nil {
				fmt.Fprintln(os.Stderr, err.Error())
				os.Exit(1)
			}

			err = ioutil.WriteFile(*output+"/blackbox/blackbox.go", []byte(blackbox), 0666)
			if err != nil {
				fmt.Println(err)
				os.Exit(1)
			}
		}
	}

	// Check if file has recursion.
	if tools.CheckRecursion(string(src), *output) {
		fmt.Println("Recursion detected : applying transformation")
		fmt.Println("maxdepth set to", *maxdepth)
		src = []byte(tools.RewriteRecursion(string(src), *output, *maxdepth, *ignorefunc != ""))
		f, err = parser.ParseFile(fset, path, src, 0)
		if err != nil {
			return "", err
		}
	}

	for _, g := range goroutines {
		g.Status = meta.StateSize()
		meta.Create("Status", g.Name)
		g.LineCounter = meta.StateSize()
		meta.Create("LineCounter", g.Name)
	}

	res = transform.ApplyAll(f, meta, []transform.Transform{
		&transform.Channel{},
		&transform.NormalizeDeclarations{},
		&transform.Alive{},
		&transform.ArrayTreat{},
		&transform.ElseFull{},
		&transform.ArithmeticCall{},
		&transform.RegularizedAffect{},
		&transform.FunctionDefs{},
		&transform.RoutineAssignments{
			Functions: goroutines,
		},
		&transform.AfterChannel{},
		&transform.LocalVariableAssignments{},
		&cfg.Transform{},
		&transform.RoutineCounter{
			Functions: goroutines,
		},
		&transform.Format{},
	})

	templ, count := "", 0
	tmpChan, tmpName, channels := channel.Channel{}, "", []channel.Channel{}
	for i, elmt := range meta.GetVariables() {
		if strings.Contains(elmt, "main.array_") {
			if strings.Contains(elmt, "_state_") {
				if templ == "" {
					templ = elmt
					tmpName = string(templ[11])
					for _, str := range templ[12:] {
						if string(str) == "_" {
							break
						} else {
							tmpName += string(str)
						}
					}
					tmpChan.State = i
					count++
				} else {
					count++
				}
			} else if strings.Contains(elmt, "_value_0") {
				tmpChan.Value = i
				tmpChan.Length = count
			}
		} else if elmt == "main.isalive_"+tmpName+"_value" {
			channels = append(channels, tmpChan)
			templ, count = "", 0
		}
	}

	// Generate all the cases of the main switch case.
	cases := make([]ast.Stmt, len(meta.GetFunctions())+1)

	// The first case it the entry point of the program, where we initialize
	// the function counter to the index of the program's main function,
	// and set its caller counter to a non-existent index, thereby ending the
	// program once it finishes executing the main function.
	cases[0] = &ast.CaseClause{
		List: []ast.Expr{
			transform.Int(0),
		},
		Body: []ast.Stmt{
			meta.Set(meta.FunctionCounter(), transform.Int(meta.MainFunction().Index)),
			// When the main function completes, break from the switch by jumping to a
			// non-existent case.
			meta.Set(meta.MainFunction().Caller, transform.Int(len(meta.GetFunctions())+1)),
			&ast.ReturnStmt{
				Results: []ast.Expr{
					meta.State(),
				},
			},
		},
	}

	// All other cases are simply a matter of calling the right function with the
	// program state.
	for _, fr := range meta.GetFunctions() {
		cases[fr.Index] = &ast.CaseClause{
			List: []ast.Expr{
				transform.Int(fr.Index),
			},
			Body: []ast.Stmt{
				&ast.ReturnStmt{
					Results: []ast.Expr{
						&ast.CallExpr{
							Fun: fr.Ident,
							Args: []ast.Expr{
								meta.State(),
							},
						},
					},
				},
			},
		}
	}

	// Creates the main function's switch case.
	switchStmt := &ast.SwitchStmt{
		Init: nil,
		Tag:  meta.Get(meta.FunctionCounter()),
		Body: &ast.BlockStmt{
			List: cases,
		},
	}

	// Produce array of values for RERS
	var rers_values = []int{}
	var rers_position = 0
	var rers_active = 0

	if *rers != "" {
		rers_active = 1
		var unspacify = regexp.MustCompile(` *`)
		s := strings.Split(unspacify.ReplaceAllString(*rers, ""), ";")

		for _, i := range s {
			j, err := strconv.Atoi(i)
			if err != nil {
				panic(err)
			}
			rers_values = append(rers_values, j)
		}
		for i, e := range meta.GetVariables() {
			if strings.Contains(e, "__RERS__") {
				rers_position = i
				break
			}
		}
	}

	// We also need to add declarations for the state size, the type of the state
	// itself, the names of the different variables of our program, and finally
	// the entry function itself.
	transform.InjectDecls(
		res.(*ast.File),
		[]ast.Decl{
			// Do we want a support for RERS?
			decl.Int(token.CONST, "G2PRersActive", rers_active),
			decl.Int(token.CONST, "G2PRersPos", rers_position),
			decl.IntArray("G2PRersValues", rers_values),

			// size of state vector
			decl.Int(token.CONST, "G2PStateSize", meta.StateSize()),

			// type declaration for state vector
			decl.Type(
				meta.StateType(),
				&ast.ArrayType{
					Len: ast.NewIdent("G2PStateSize"),
					Elt: ast.NewIdent("int"),
				},
			),

			// array of strings containing variable names for all variables contained
			// in the state vector
			decl.Array(
				"G2PVariableNames",
				&ast.ArrayType{
					Len: &ast.BasicLit{
						Kind:  token.INT,
						Value: strconv.Itoa(meta.StateSize()),
					},
					Elt: &ast.StarExpr{
						X: &ast.SelectorExpr{
							X:   ast.NewIdent("C"),
							Sel: ast.NewIdent("char"),
						},
					},
				},
				variableNamesToC(meta.GetVariables()),
			),

			decl.GoroutineArray("G2PGoroutines", goroutines),
			decl.ChannelArray("G2PChannels", channels),

			// entrypoint to get a successor from a state vector
			decl.Function(
				"G2PMain",
				[]*ast.Field{ // parameters
					&ast.Field{
						Names: []*ast.Ident{meta.State()},
						Type:  meta.StateType(),
					},
				},
				[]*ast.Field{ // return values
					&ast.Field{
						//Type: transform.ArrayOf(meta.StateType()),
						Type: meta.StateType(),
					},
				},
				[]ast.Stmt{ // function body
					switchStmt,
					&ast.ReturnStmt{
						Results: []ast.Expr{
							meta.State(),
						},
					},
				},
			),
		},
	)
	astutil.AddImport(fset, res.(*ast.File), "C")
	return f.Name.Name, format.Node(file, fset, res)
}

// Copies files over and format them according to the template data.
func copyFile(src string, dst string, data interface{}) error {
	srcBytes, err := ioutil.ReadFile(src)
	if err != nil {
		return err
	}

	dstFile, err := os.Create(dst)
	if err != nil {
		return err
	}
	defer dstFile.Close()

	t, err := template.New("temp").Parse(string(srcBytes))
	if err != nil {
		return err
	}

	return t.Execute(dstFile, data)
}

func duplicateGoroutines(filepath string) (string, []*goroutine.Goroutine) {
	fset := token.NewFileSet()

	node, err := parser.ParseFile(fset, filepath, nil, 0)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

	tmpFile, err := ioutil.TempFile("", "*_go2pins.go")
	if err != nil {
		panic("couldn't create go2pins temp file")
	}
	defer tmpFile.Close()

	srcFile, err := os.Open(filepath)
	if err != nil {
		panic("couldn't open given file")
	}
	defer srcFile.Close()

	_, err = io.Copy(tmpFile, srcFile)
	if err != nil {
		panic("couldn't copy file contents")
	}

	goroutines := make([]*goroutine.Goroutine, 0)
	gofuncs := make([]string, 0)
	ast.Inspect(node, func(n ast.Node) bool {
		gostmt, ok := n.(*ast.GoStmt)
		if !ok {
			return true
		}

		// get line of go call
		absPos := fset.PositionFor(gostmt.Pos(), false)

		chanToAdd := ""
		var fun *ast.Ident
		switch f := gostmt.Call.Fun.(type) {
		case *ast.Ident:
			fun = f
			switch decl := f.Obj.Decl.(type) {
			case *ast.FuncDecl:
				counter := 0
				for _, elmt := range decl.Type.Params.List {
					switch elmt.Type.(type) {
					case *ast.ChanType:
						for j, chanParamName := range elmt.Names {
							switch chanName := gostmt.Call.Args[counter+j].(type) {
							case *ast.Ident:
								chanToAdd += "_" + chanName.Name
							default:
								chanToAdd += "_" + chanParamName.Name
							}
						}
					}
					counter += len(elmt.Names)
				}
			}
		case *ast.SelectorExpr:
			fun = f.Sel
		default:
			panic("unsupported go call")
		}

		f, ok := fun.Obj.Decl.(*ast.FuncDecl)
		if !ok {
			panic("object should be a function")
		}

		fmt.Fprintln(tmpFile) // skip new line
		tmpName := f.Name
		if *ignorefunc != "" {
			for _, elmt := range toIgnore {
				if elmt == tmpName.Name {
					fmt.Println("WARNING: Blackbox can not be applied to goroutine call.")
					fmt.Println("\tFunc: " + elmt)
				}
			}
		}
		gofuncs = append(gofuncs, f.Name.Name)
		name := "_go_" + f.Name.Name + chanToAdd + "_" + strconv.Itoa(absPos.Line)
		f.Name = ast.NewIdent(name)
		format.Node(tmpFile, fset, f)
		f.Name = tmpName
		goroutines = append(goroutines, &goroutine.Goroutine{
			Name: name,
		})

		return true
	})

	input, err := ioutil.ReadFile(tmpFile.Name())
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	newInput := string(input)

	for i, gf := range gofuncs {
		newInput = strings.Replace(newInput, "go "+gf+"(", "go "+goroutines[i].Name+"(", 1)
	}

	err = ioutil.WriteFile(tmpFile.Name(), []byte(newInput), 0666)
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	return tmpFile.Name(), goroutines
}

type packageTemplate struct {
	PackagePath string
	PackageName string
}

var toIgnore []string

var (
	astprint       = flag.Bool("print-ast", false, "Print AST of <file>")
	callgraph      = flag.Bool("callgraph", false, "Display callgraph of <file>")
	forceOverwrite = flag.Bool("f", false, "overwrite output file if it exists")
	ignorefunc     = flag.String("ignore-fn", "", "Function will be considered as non-existent\n./go2pins -ignore-fn \"func1;func2;func3\" <file>")
	maxdepth       = flag.Int("maxdepth", 10, "fix the maximum depth of recursion for a recursive file")
	output         = flag.String("o", "output", "output `directory`")
	rers           = flag.String("rers", "", "support for RERS challenge: instruction \"__RERS += __RERS__ + 1\" will be\nreplaced to simulate environment, i.e. this instruction will produce as many\nsuccessors as values in parameters.\n./go2pins -rers \"val1;val2;val3\"")
)

func init() {
	flag.Usage = func() {
		fmt.Fprintln(os.Stderr, "Usage: "+os.Args[0]+" [options] <file>\n")
		flag.PrintDefaults()
	}
}

func main() {
	flag.Parse()
	if len(flag.Args()) != 1 {
		flag.Usage()
		os.Exit(1)
	}

	if *astprint {
		tools.PrintAst(flag.Args()[0])
	} else if *callgraph {
		tools.PrintCG(tools.CallGraph(flag.Args()[0]))
	} else {
		relPackagePath := flag.Args()[0]
		basePackagePath := path.Base(relPackagePath)

		// Check whether the output/ directory already exists.
		if _, err := os.Stat(*output); err == nil {
			if !*forceOverwrite {
				fmt.Fprintln(os.Stderr, *output+" already exists")
				os.Exit(1)
			}
			err = os.RemoveAll(*output)
			if err != nil {
				fmt.Fprintln(os.Stderr, "couldn't delete "+*output)
				os.Exit(1)
			}
		}

		// Create the output directory.
		if err := os.MkdirAll(*output, 0777); err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
			os.Exit(1)
		}

		// Create a file in output/ with the same name as the input package.
		f, err := os.Create(path.Join(*output, basePackagePath))
		if err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
			os.Exit(1)
		}
		defer f.Close()

		if *ignorefunc != "" {
			var unspacify = regexp.MustCompile(` *`)
			s := unspacify.ReplaceAllString(*ignorefunc, "")
			toIgnore = strings.Split(s, ";")
		}

		tmpFilePath, goroutines := duplicateGoroutines(relPackagePath)
		// FIXME: do we want this ?
		// defer os.Remove(tmpFilePath)
		flag.Args()[0] = tmpFilePath

		src, err := ioutil.ReadFile(tmpFilePath)
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}

		// Compile and write the package to the output file.
		packageName, err := compileTo(basePackagePath, src, f, goroutines, rers)
		if err != nil {
			fmt.Fprintln(os.Stderr, err.Error())
			os.Exit(1)
		}

		tools.FormatArray(*output+"/"+basePackagePath, "var G2PVariableNames =")

		// Get project root directory to locate boilerplate/ dir
		_, currentPath, _, ok := runtime.Caller(0)
		if !ok {
			fmt.Fprintln(os.Stderr, "can't get path of current file")
			os.Exit(1)
		}
		go2pinsDir := path.Dir(currentPath)

		boilerplateFiles := []string{
			"Makefile",
			"go2pins.c",
			"go2pins.h",
			"main.go",
		}
		data := packageTemplate{basePackagePath, packageName}
		for _, filename := range boilerplateFiles {
			from := path.Join(go2pinsDir, "boilerplate/"+filename)
			to := path.Join(*output, filename)
			copyFile(from, to, data)
		}

		fmt.Println("Preprocessing done.\n\nTo display the state space of you program, follows the instructions below")
		fmt.Println("      cd", *output)
		fmt.Println("      make")
		fmt.Println("      <path_to_spot>/ltsmin/modelcheck -gm main.dve2C 'true'")
	}
}