GO AST based transformations


(Uwimbabazi) #1

Dear Gophers,

I have been facing an issue of adding elements of type ast.IndexExpr to an empty list at once.
For example in the following code. When I have: Put(CreateTuple(name, 10), A), I wish to get the contents of A such that I could have at the transformed output program something like: A[i] = uri[s1], uri[s2], uri[s3] rather than having only the last element: A[0] = uri[s3]. The variable name of uri[s1], uri[s2], uri[s3] is indexpr02. Then, when I place it to the Rhs of assignment06. I ony get the last elements A[i] = uri[s3]. Any help is really appreciated.
// This is the program that transforms the Put action of the given input program,
// by inserting the transformed Put actions in the output program: e.g…

// spaceid.Put(tuple)

// into :

// Put(CreateTuple(tuple), a)
// WHERE: a represents the spaces

// Transform GetP and QueryP actions
// s1.GetP(&description, &key) into: GetP(“A”, &key, uri[s1], uri[s2])
////…

//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

package main

// import useful packages
import (
“flag”
“fmt”
“go/ast”
“go/parser”
“go/printer”
“go/token”
“log”
“os”
“reflect”
)

func translation3(filename string) {

//////////
////////// Section 1: (parsing)create the AST from the input program
//////////

//Create a new FileSet that represents a set of source files for the parser.
fset := token.NewFileSet()
// Parse the source filename and comments and add them to AST.
node, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
if err != nil {
	log.Fatal(err)
}

fmt.Printf(" ====> INPUT PROGRAM: ===== \n")
//Print the input program file.
printer.Fprint(os.Stdout, fset, node)

//////////
////////// Section 2: / transformation of the AST
//////////

// Generate useful identifiers
identurl := ast.NewIdent("uri")
identm := ast.NewIdent("m")
var basiclit1 *ast.BasicLit
// generate: a
identa := ast.NewIdent("a")

// Declarations of useful statements and expressions.
//var indexpr02 *ast.IndexExpr
var indexpr02 *ast.IndexExpr
var callexpr01 *ast.CallExpr
var assignmentO2 *ast.AssignStmt
var assignmentO1 *ast.AssignStmt
var assignment04 *ast.AssignStmt
// Declare useful statements for the Sub-ASTs prepared in order to be used globally
var assignment05 *ast.AssignStmt
var assignment06 *ast.AssignStmt

// target locations for put operations // this is an empty array
targets := []ast.Expr{}

//  Traversing through the AST starting from all declarations
for _, f := range node.Decls {
	var list []ast.Stmt

	// Find functions
	fn, ok := f.(*ast.FuncDecl)
	if !ok {
		continue
	}

	// Find assignment statements in the body list
	for _, k := range fn.Body.List {
		switch reflect.TypeOf(k).String() {
		case "*ast.AssignStmt":
			astmt, ok := k.(*ast.AssignStmt)
			if ok {
				switch reflect.TypeOf(astmt.Rhs[0]).String() {
				case "*ast.CallExpr":
					// Invoke the function NewSpace to extract all the spaces created
					if ok && astmt.Rhs[0].(*ast.CallExpr).Fun.(*ast.Ident).Name == "NewSpace" && reflect.TypeOf(astmt.Lhs[0]).String() == "*ast.Ident" {

						// Extract left hand side of assignment statement
						// (i.e., the identifier of the Space object).
						// Extract first argument in the function call at the right hand side of the assignment statement
						// (i.e., the URI of the space).

						spaceid := astmt.Lhs[0].(*ast.Ident)
						spaceuri := astmt.Rhs[0].(*ast.CallExpr).Args[0].(*ast.BasicLit)

						// Prepare sub-ASTs to inject in the target program
						// i.e., &spaceid, uri[spaceid], and m[uri]
						space := &ast.UnaryExpr{Op: token.AND, X: spaceid}      //  &spaceid
						indexpr01 := &ast.IndexExpr{X: identm, Index: spaceuri} // m[uri]
						indexpr02 = &ast.IndexExpr{X: identurl, Index: spaceid} // uri[spaceid]

						// Add assignment statements for m[] and uri[]
						// transform spaceuri  ->   m[spaceuri]
						assignmentO1 = &ast.AssignStmt{Lhs: []ast.Expr{indexpr01}, Tok: token.ASSIGN, Rhs: []ast.Expr{space}}
						assignmentO2 = &ast.AssignStmt{Lhs: []ast.Expr{indexpr02}, Tok: token.ASSIGN, Rhs: []ast.Expr{astmt.Rhs[0].(*ast.CallExpr).Args[0].(*ast.BasicLit)}}

						// Join the two parts (assignement01, assignement02) above:
						// m[spaceuri]  ->  &spaceid
						// uri[spaceid] ->  spaceuri
						list = append(list, assignmentO1, assignmentO2)
						targets = append(targets, indexpr02)
						//targets = append(targets, Name)

					}

				default:
					list = append(list, k)
				}
			}
		default:
			list = append(list, k)
		}

		fn.Body.List = list

	}
}

// ...Traversing through the block of statements
ast.Inspect(node, func(n ast.Node) bool {
	fn, ok := n.(*ast.BlockStmt)
	if ok {
		var expr ast.Expr

		// Find in the list, the expressions statements
		for _, k := range fn.List {
			if reflect.TypeOf(k).String() == "*ast.ExprStmt" {
				fn1, ok := k.(*ast.ExprStmt)

				// Invoke the function Put to extract all the arguments associated with it
				//...
				if ok && reflect.TypeOf(fn1.X.(*ast.CallExpr).Fun).String() == "*ast.SelectorExpr" {
					if fn1.X.(*ast.CallExpr).Fun.(*ast.SelectorExpr).Sel.Name == "Put" {
						// Extract all the args associated with the function call "Put" and create a new args
						newArgs := make([]ast.Expr, len(fn1.X.(*ast.CallExpr).Args))
						copy(newArgs, fn1.X.(*ast.CallExpr).Args)

						// Prepare the sub-ASTs for the: a := make([]string, 2)
						//....

						// generate: 2
						basiclit := &ast.BasicLit{Kind: token.STRING, Value: "10"}
						// generate: string
						identstring := ast.NewIdent("string")
						// generate: make
						identmake := ast.NewIdent("make")

						// generate: []string
						arraytype := &ast.ArrayType{Elt: identstring}
						// generate: make([]string, 2)
						callexpr02 := &ast.CallExpr{Fun: identmake, Args: []ast.Expr{arraytype, basiclit}}
						//generate: a := make([]string, 2)
						assignment05 = &ast.AssignStmt{Lhs: []ast.Expr{identa}, Tok: token.DEFINE, Rhs: []ast.Expr{callexpr02}}

						basiclit1 = &ast.BasicLit{Kind: token.STRING, Value: "\"uri[s1]\""}

						//basiclit1 = &ast.BasicLit{Kind: token.STRING, Value: "" + Name + ""}
						// generate: 0
						basiclit2 := &ast.BasicLit{Kind: token.STRING, Value: "0"}
						// generate: a[0]
						indexexpr03 := &ast.IndexExpr{X: identa, Index: basiclit2}
						// generate: a[0] := "uri[s1]"

						assignment06 = &ast.AssignStmt{Lhs: []ast.Expr{indexexpr03}, Tok: token.ASSIGN, Rhs: []ast.Expr{indexpr02}}

						// Prepare the sub-ASTs : t := CreateTuple(args)
						//....

						// generate : CreateTuple
						identcreatetuple := ast.NewIdent("CreateTuple")
						// generate: t
						identt := ast.NewIdent("t")
						// generate: CreateTuple(newargs)
						callexpr01 = &ast.CallExpr{Fun: identcreatetuple, Args: newArgs}
						// generate: t := CreateTuple(newargs)
						assignment04 = &ast.AssignStmt{Lhs: []ast.Expr{identt}, Tok: token.DEFINE, Rhs: []ast.Expr{callexpr01, indexpr02}}

						// For each element in targets that is not nil: get new list of args
						for i := range targets {
							if targets[i] != nil {
								newArgs = append(newArgs, targets[i])

							}
						}

						//Generate: Put(CreateTuple(args), a)
						expr = &ast.CallExpr{Fun: fn1.X.(*ast.CallExpr).Fun.(*ast.SelectorExpr).Sel, Args: []ast.Expr{callexpr01, identa}}

					}

					if expr != nil {
						fn1.X = expr
					}
				}
			}
		}
	}
	return true
})

// Traversing through the block of declarations
for _, f := range node.Decls {
	// declare an array statements
	var list2 []ast.Stmt
	// combine the following prepared sub-ASTs:
	//a := make([]string, 2)
	//a[0] := uri[s1]
	//a[1] := uri[s2]
	//...

	list2 = append(list2, assignment05, assignment06)

	// Find the functions
	fn, ok := f.(*ast.FuncDecl)
	if !ok {
		continue
	}

	//Put(CreateTuple(tuple), a)
	//....

	if fn.Name.Name != "main" {
		fn.Body.List = append(list2, fn.Body.List...)

	}
}

//////////
////////// Section 3: Pretty print the AST,
//////////            in this way we obtain the modified program.
//////////

fmt.Println(" ~~~~~~~TRANSFORMED OUTPUT PROGRAM ~~~~~~~~~ ")
// Print the transformed output program
printer.Fprint(os.Stdout, fset, node)

}

//////////
////////// Section 4: Make sure to xecute the inputfilename using the command lines.
//////////

// Execute the inputfilename in the command lines
var inputfilename *string

func usage() {
fmt.Fprintf(os.Stderr, “usage: myprog -i [inputfile]\n”)
flag.PrintDefaults()
os.Exit(2)
}

func initparams() {
inputfilename = flag.String(“i”, “”, “input filename”)
flag.Parse()

if *inputfilename == "" {
	usage()
}

}

func main() {
translation3(“testfiles/input3b.go”)

}

~~~~~~~TRANSFORMED OUTPUT PROGRAM ~~~~~~~~~
package main

import (
. “github.com/pspaces/gospace
)

func main() {
m[“tcp://host:192100/s1”] = &s1
uri[s1] = “tcp://host:192100/s1”
m[“tcp://host:192101/s2”] = &s2
uri[s2] = “tcp://host:192101/s2”
m[“tcp://host:13456/s3”] = &s3
uri[s3] = “tcp://host:13456/s3”
go Process1(&s1)
go Process1(&s2)
go Process1(&s3)
}

func Process1(s1 *Space) {
a := make([]string, 10)
a[0] = uri[s3]

var name bool
Put(CreateTuple("A", 10), a)
Put(CreateTuple(name, 10), a)

}