recursion.go

// Copyright (C) 2020 Laboratoire de Recherche et Developpement
// de l'EPITA (LRDE).
//
// This file is part of Go2Pins, a tool for Golang model-checking
//
// Go2Pins is free software; you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 3 of the License, or
// (at your option) any later version.
//
// Go2Pins is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
// or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
// License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package tools

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/parser"
	"go/printer"
	"go/token"
	"log"
	"strconv"
)

type FuncInfo struct {
	Name                  string
	IsRecursive           bool
	TransitivelyRecursive bool
	Calls                 []string
}

func collectFuncInfo(node *ast.FuncDecl) *FuncInfo {
	// Check wether the function is recursive
	var is_recursive = false
	calls := []string{}
	ast.Inspect(node, func(n ast.Node) bool {
		callexp, ok := n.(*ast.CallExpr)
		if !ok {
			return true
		}

		switch f := callexp.Fun.(type) {
		case *ast.Ident:
			// FIXME this is a cast operation
			// only cast to int is currently supported
			if f.Name == "int" {
				return true
			}

			if node.Name.Name == f.Name {
				is_recursive = true
			}

			ok := false
			for _, val := range calls {
				if val == f.Name {
					ok = true
					break
				}
			}
			if !ok {
				calls = append(calls, f.Name)
			}
		}

		return true
	})

	res := FuncInfo{node.Name.Name, is_recursive, false /* Unknown at this point*/, calls}
	return &res
}

func BuildCallgraph(inputfile string) map[string]*FuncInfo {
	fset := token.NewFileSet()

	node, err := parser.ParseFile(fset, "src.go", inputfile, 0)
	if err != nil {
		log.Fatal(err)
	}

	callgraph := make(map[string]*FuncInfo)
	// Walk the file to compute relevant information about recursion
	ast.Inspect(node, func(n ast.Node) bool {
		// check if we're on a `go` node
		decl, ok := n.(*ast.FuncDecl)
		if !ok {
			return true
		}

		// FIXME HANDLE GOROUTINE + LAMBDA Function

		// get position in file
		callgraph[decl.Name.Name] = collectFuncInfo(decl)
		return true
	})

	// Fill the transitively recursive field
	for key, element := range callgraph {
		visited := make(map[string]bool)
		trans_rec := false
		var indirect func(string, string)
		indirect = func(from string, lookup string) {
			visited[from] = true
			for _, call := range callgraph[from].Calls {
				if lookup == call {
					trans_rec = true
				}
				_, ok := visited[call]
				if !ok {
					indirect(call, lookup)
				}
			}
		}
		indirect(element.Name, element.Name)
		callgraph[key].TransitivelyRecursive = trans_rec
	}
	return callgraph
}

func Callgraph(inputfile string) {
	callgraph := BuildCallgraph(inputfile)
	digraph := "digraph callgraph {\n"
	for key, element := range callgraph {
		for _, call := range element.Calls {
			digraph += "  " + key + " -> " + call + ";\n"
		}
	}
	digraph += "}\n"
	fmt.Println(digraph)
}

func parseFunc(filename, functionname string) (fun *ast.FuncDecl, fset *token.FileSet) {
	fset = token.NewFileSet()
	if file, err := parser.ParseFile(fset, "src.go", filename, 0); err == nil {
		for _, d := range file.Decls {
			if f, ok := d.(*ast.FuncDecl); ok && f.Name.Name == functionname {
				fun = f
				return
			}
		}
	}
	panic("function not found " + functionname)
}

func RewriteRecursion(inputfile string, unroll_level int) (string, bool) {
	fset := token.NewFileSet()

	node, err := parser.ParseFile(fset, "src.go", inputfile, 0)
	if err != nil {
		log.Fatal(err)
	}

	result := ""
	rec := false
	for key, element := range BuildCallgraph(inputfile) {

		if !(element.IsRecursive || element.TransitivelyRecursive) {
			continue
		}

		rec = true

		// FIXME modify internal calls

		for i := 1; i <= unroll_level; i++ {
			funcAST, fset := parseFunc(inputfile, key)
			funcAST.Name.Name = funcAST.Name.Name + "_" +
				strconv.Itoa(i)

			if i != unroll_level {
				ast.Inspect(funcAST.Body, func(n ast.Node) bool {
					callexp, ok := n.(*ast.CallExpr)
					if !ok {
						return true
					}

					switch f := callexp.Fun.(type) {
					case *ast.Ident:
						f.Name = f.Name + "_" + strconv.Itoa(i+1)
					}

					return true
				})
			} else {
				funcAST.Body.List = []ast.Stmt{}

				funcAST.Body.List = append(funcAST.Body.List,
					&ast.ExprStmt{
						X: &ast.CallExpr{
							Fun: ast.NewIdent("panic"),
							Args: []ast.Expr{
								&ast.BasicLit{
									Kind:  token.STRING,
									Value: "\"Max Depth Reached\"",
								},
							},
						},
					})
			}

			var buf bytes.Buffer
			printer.Fprint(&buf, fset, funcAST)
			result = result + "\n" + buf.String() + "\n"
		}

		ast.Inspect(node, func(n ast.Node) bool {
			// check if we're on a `go` node
			decl, ok := n.(*ast.FuncDecl)
			if !ok {
				return true
			}
			if decl.Name.Name != key {
				return true
			}
			ast.Inspect(decl.Body, func(n ast.Node) bool {
				callexp, ok := n.(*ast.CallExpr)
				if !ok {
					return true
				}

				switch f := callexp.Fun.(type) {
				case *ast.Ident:
					f.Name = f.Name + "_1"
				}

				return true
			})

			return true
		})

	}

	var buf bytes.Buffer
	printer.Fprint(&buf, fset, node)
	result = buf.String() + result

	return result, rec
}