// addCode searches for main func in data, and updates AST code // adding tracing functions. func addCode(path string) ([]byte, error) { var conf loader.Config if _, err := conf.FromArgs([]string{path}, false); err != nil { return nil, err } prog, err := conf.Load() if err != nil { return nil, err } // check if runtime/trace already imported for i, _ := range prog.Imported { if i == "runtime/trace" { return nil, ErrImported } } pkg := prog.Created[0] // TODO: find file with main func inside astFile := pkg.Files[0] // add imports astutil.AddImport(prog.Fset, astFile, "os") astutil.AddImport(prog.Fset, astFile, "runtime/trace") astutil.AddImport(prog.Fset, astFile, "time") // add start/stop code ast.Inspect(astFile, func(n ast.Node) bool { switch x := n.(type) { case *ast.FuncDecl: // find 'main' function if x.Name.Name == "main" && x.Recv == nil { stmts := createTraceStmts() stmts = append(stmts, x.Body.List...) x.Body.List = stmts return true } } return true }) var buf bytes.Buffer err = printer.Fprint(&buf, prog.Fset, astFile) if err != nil { return nil, err } return buf.Bytes(), nil }
func args(filename string) { b, err := ioutil.ReadFile(filename) if err != nil { panic(err) } p := parser.ParseFromString(filename, string(b)+"\n") a := generator.GenerateAST(p) fset := token.NewFileSet() defaultImports := []string{"github.com/gsp-lang/stdlib/prelude", "github.com/gsp-lang/gsp/core"} for _, defaultImport := range defaultImports { split := strings.Split(defaultImport, "/") pkgName := split[len(split)-1] if !(a.Name.Name == "prelude" && pkgName == "prelude") { astutil.AddImport(fset, a, defaultImport) } } if a.Name.Name != "prelude" { a.Decls = append(a.Decls, &ast.GenDecl{ Tok: token.VAR, Specs: []ast.Spec{&ast.ValueSpec{ Names: []*ast.Ident{&ast.Ident{Name: "_"}}, Values: []ast.Expr{&ast.Ident{Name: "prelude.Len"}}, }}, }) } var buf bytes.Buffer printer.Fprint(&buf, fset, a) fmt.Printf("%s\n", buf.String()) }
func args(filename string) { b, err := ioutil.ReadFile(filename) if err != nil { panic(err) } p := parser.ParseFromString(filename, string(b)+"\n") a := generator.GenerateAST(p) fset := token.NewFileSet() defaultImports := []string{"github.com/gsp-lang/stdlib/prelude", "github.com/gsp-lang/gsp/core"} for _, defaultImport := range defaultImports { split := strings.Split(defaultImport, "/") pkgName := split[len(split)-1] if !(a.Name.Name == "prelude" && pkgName == "prelude") { if pkgName == "prelude" { astutil.AddNamedImport(fset, a, "_", defaultImport) } else { astutil.AddImport(fset, a, defaultImport) } } } var buf bytes.Buffer printer.Fprint(&buf, fset, a) fmt.Printf("%s\n", buf.String()) }
// Inline replaces each instance of identifier k with v.Ident in ast.File f, // for k, v := range m. // For all inlines that were triggeres it also adds imports from v.Imports to f. // In addition, it removes top level type declarations of the form // type k ... // for all k in m. // // Every k in m should be a valid identifier. // Every v.Ident should be a valid expression. func Inline(fset *token.FileSet, f *ast.File, m map[string]Target) error { // Build the inline map. im := map[string]reflect.Value{} for k, v := range m { expr, err := parser.ParseExpr(k) if err != nil { return fmt.Errorf("failed to parse `%s`: %s", k, err) } if _, ok := expr.(*ast.Ident); !ok { return fmt.Errorf("expected identifier, got %s which is %T", k, expr) } expr, err = parser.ParseExpr(v.Ident) if err != nil { return fmt.Errorf("failed to parse `%s`: %s", v.Ident, err) } s := v.Ident if _, ok := expr.(*ast.StarExpr); ok { s = fmt.Sprintf("(%s)", s) } im[k] = reflect.ValueOf(ast.Ident{Name: s}) } // Filter `type XXX ...` declarations out if we are inlining XXX. cmap := ast.NewCommentMap(fset, f, f.Comments) to := 0 for _, d := range f.Decls { skip := false if t, ok := d.(*ast.GenDecl); ok { for _, s := range t.Specs { ts, ok := s.(*ast.TypeSpec) if !ok { continue } if _, ok = im[ts.Name.String()]; ok { skip = true } } } if !skip { f.Decls[to] = d to++ } } if to != len(f.Decls) { f.Decls = f.Decls[:to] // Remove comments for the declarations that were filtered out. f.Comments = cmap.Filter(f).Comments() } // Add imports for the inlines that were triggered. for k := range inline(im, f) { for _, imp := range m[k].Imports { astutil.AddImport(fset, f, imp) } } return nil }
func actionImport(s *Session, arg string) error { if arg == "" { return fmt.Errorf("arg required") } path := strings.Trim(arg, `"`) // check if the package specified by path is importable _, err := s.Types.Importer.Import(path) if err != nil { return err } astutil.AddImport(s.Fset, s.File, path) return nil }
func addImports(file *ast.File, fset *token.FileSet, dirPath string) (*ast.File, *token.FileSet, error) { imports, err := getImports(dirPath, fset) if err != nil { return nil, nil, err } for _, s := range imports { unquotedPath, err := strconv.Unquote(s.Path.Value) if err != nil { return nil, nil, err } if s.Name != nil { astutil.AddNamedImport(fset, file, s.Name.Name, unquotedPath) continue } astutil.AddImport(fset, file, unquotedPath) } return file, fset, nil }
func fixImports(fset *token.FileSet, f *ast.File, filename string) (added []string, err error) { // refs are a set of possible package references currently unsatisfied by imports. // first key: either base package (e.g. "fmt") or renamed package // second key: referenced package symbol (e.g. "Println") refs := make(map[string]map[string]bool) // decls are the current package imports. key is base package or renamed package. decls := make(map[string]*ast.ImportSpec) abs, err := filepath.Abs(filename) if err != nil { return nil, err } srcDir := path.Dir(abs) // collect potential uses of packages. var visitor visitFn visitor = visitFn(func(node ast.Node) ast.Visitor { if node == nil { return visitor } switch v := node.(type) { case *ast.ImportSpec: if v.Name != nil { decls[v.Name.Name] = v } else { local := importPathToName(strings.Trim(v.Path.Value, `\"`), srcDir) decls[local] = v } case *ast.SelectorExpr: xident, ok := v.X.(*ast.Ident) if !ok { break } if xident.Obj != nil { // if the parser can resolve it, it's not a package ref break } pkgName := xident.Name if refs[pkgName] == nil { refs[pkgName] = make(map[string]bool) } if decls[pkgName] == nil { refs[pkgName][v.Sel.Name] = true } } return visitor }) ast.Walk(visitor, f) // Nil out any unused ImportSpecs, to be removed in following passes unusedImport := map[string]string{} for pkg, is := range decls { if refs[pkg] == nil && pkg != "_" && pkg != "." { name := "" if is.Name != nil { name = is.Name.Name } unusedImport[strings.Trim(is.Path.Value, `"`)] = name } } for ipath, name := range unusedImport { if ipath == "C" { // Don't remove cgo stuff. continue } astutil.DeleteNamedImport(fset, f, name, ipath) } // Search for imports matching potential package references. searches := 0 type result struct { ipath string name string err error } results := make(chan result) for pkgName, symbols := range refs { if len(symbols) == 0 { continue // skip over packages already imported } go func(pkgName string, symbols map[string]bool) { ipath, rename, err := findImport(pkgName, symbols, filename) r := result{ipath: ipath, err: err} if rename { r.name = pkgName } results <- r }(pkgName, symbols) searches++ } for i := 0; i < searches; i++ { result := <-results if result.err != nil { return nil, result.err } if result.ipath != "" { if result.name != "" { astutil.AddNamedImport(fset, f, result.name, result.ipath) } else { astutil.AddImport(fset, f, result.ipath) } added = append(added, result.ipath) } } return added, nil }
// Transform applies the transformation to the specified parsed file, // whose type information is supplied in info, and returns the number // of replacements that were made. // // It mutates the AST in place (the identity of the root node is // unchanged), and may add nodes for which no type information is // available in info. // // Derived from rewriteFile in $GOROOT/src/cmd/gofmt/rewrite.go. // func (tr *Transformer) Transform(info *types.Info, pkg *types.Package, file *ast.File) int { if !tr.seenInfos[info] { tr.seenInfos[info] = true mergeTypeInfo(tr.info, info) } tr.currentPkg = pkg tr.nsubsts = 0 if tr.verbose { fmt.Fprintf(os.Stderr, "before: %s\n", astString(tr.fset, tr.before)) fmt.Fprintf(os.Stderr, "after: %s\n", astString(tr.fset, tr.after)) } var f func(rv reflect.Value) reflect.Value f = func(rv reflect.Value) reflect.Value { // don't bother if val is invalid to start with if !rv.IsValid() { return reflect.Value{} } rv = apply(f, rv) e := rvToExpr(rv) if e != nil { savedEnv := tr.env tr.env = make(map[string]ast.Expr) // inefficient! Use a slice of k/v pairs if tr.matchExpr(tr.before, e) { if tr.verbose { fmt.Fprintf(os.Stderr, "%s matches %s", astString(tr.fset, tr.before), astString(tr.fset, e)) if len(tr.env) > 0 { fmt.Fprintf(os.Stderr, " with:") for name, ast := range tr.env { fmt.Fprintf(os.Stderr, " %s->%s", name, astString(tr.fset, ast)) } } fmt.Fprintf(os.Stderr, "\n") } tr.nsubsts++ // Clone the replacement tree, performing parameter substitution. // We update all positions to n.Pos() to aid comment placement. rv = tr.subst(tr.env, reflect.ValueOf(tr.after), reflect.ValueOf(e.Pos())) } tr.env = savedEnv } return rv } file2 := apply(f, reflect.ValueOf(file)).Interface().(*ast.File) // By construction, the root node is unchanged. if file != file2 { panic("BUG") } // Add any necessary imports. // TODO(adonovan): remove no-longer needed imports too. if tr.nsubsts > 0 { pkgs := make(map[string]*types.Package) for obj := range tr.importedObjs { pkgs[obj.Pkg().Path()] = obj.Pkg() } for _, imp := range file.Imports { path, _ := strconv.Unquote(imp.Path.Value) delete(pkgs, path) } delete(pkgs, pkg.Path()) // don't import self // NB: AddImport may completely replace the AST! // It thus renders info and tr.info no longer relevant to file. var paths []string for path := range pkgs { paths = append(paths, path) } sort.Strings(paths) for _, path := range paths { astutil.AddImport(tr.fset, file, path) } } tr.currentPkg = nil return tr.nsubsts }
// ProcessFileAST processes the files using golang's AST parser func ProcessFileAST(filePath string, from string, to string) { //Colors to be used on the console red := ansi.ColorCode("red+bh") white := ansi.ColorCode("white+bh") yellow := ansi.ColorCode("yellow+bh") blackOnWhite := ansi.ColorCode("black+b:white+h") //Reset the color reset := ansi.ColorCode("reset") fmt.Println(blackOnWhite+"Processing file", filePath, "in SAFE MODE", reset) // New FileSet to parse the go file to fSet := token.NewFileSet() // Parse the file file, err := parser.ParseFile(fSet, filePath, nil, 0) if err != nil { fmt.Println(err) } // Get the list of imports from the ast imports := astutil.Imports(fSet, file) // Keep track of number of changes numChanges := 0 // Iterate through the imports array for _, mPackage := range imports { for _, mImport := range mPackage { // Since astutil returns the path string with quotes, remove those importString := strings.TrimSuffix(strings.TrimPrefix(mImport.Path.Value, "\""), "\"") // If the path matches the oldpath, replace it with the new one if strings.Contains(importString, from) { //If it needs to be replaced, increase numChanges so we can write the file later numChanges++ // Join the path of the import package with the remainder from the old one after removing the old import package replacePackage := strings.Replace(importString, from, to, -1) fmt.Println(red + "Updating import " + reset + white + importString + reset + red + " to " + reset + white + replacePackage + reset) // Remove the old import and replace it with the replacement astutil.DeleteImport(fSet, file, importString) astutil.AddImport(fSet, file, replacePackage) } } } // If the number of changes are more than 0, write file if numChanges > 0 { // Print the new AST tree to a new output buffer var outputBuffer bytes.Buffer printer.Fprint(&outputBuffer, fSet, file) ioutil.WriteFile(filePath, outputBuffer.Bytes(), os.ModePerm) fmt.Println(yellow+ "File", filePath, "saved after", numChanges, "changes", reset, "\n\n") } else { fmt.Println(yellow+ "No changes to write on this file.", reset, "\n\n") } }
func (d *DirectiveList) Save() error { fset := token.NewFileSet() f, err := parser.ParseFile(fset, d.file, nil, 0) if err != nil { return ErrInvalidDirectiveFile } imps := astutil.Imports(fset, f) for _, imp := range imps[0] { name := "" if imp.Name != nil && imp.Name.String() != "." { name = imp.Name.String() } path, _ := strconv.Unquote(imp.Path.Value) if name != "" { astutil.DeleteNamedImport(fset, f, name, path) } else { astutil.DeleteImport(fset, f, path) } } astutil.AddImport(fset, f, "github.com/mholt/caddy/caddy/https") astutil.AddImport(fset, f, "github.com/mholt/caddy/caddy/parse") astutil.AddImport(fset, f, "github.com/mholt/caddy/caddy/setup") astutil.AddImport(fset, f, "github.com/mholt/caddy/middleware") importName := "" for _, dir := range d.list { if dir.Removed { continue } if dir.Core == false && len(dir.ImportPath) != 0 { importName = strings.Replace(dir.Name, "-", "_", -1) astutil.AddImport(fset, f, fmt.Sprintf("{{import-%s}}", importName)) } } var buf bytes.Buffer err = printer.Fprint(&buf, fset, f) if err != nil { return err } out := buf.String() f, err = parser.ParseFile(token.NewFileSet(), "", out, 0) node, ok := f.Scope.Lookup("directiveOrder").Decl.(ast.Node) if !ok { return ErrInvalidDirectiveFile } begin := out[0 : node.Pos()-1] end := out[node.End():len(out)] dirOrder := "" for _, dir := range d.list { if dir.Removed { continue } comment := "" if dir.Active == false { comment = "//@caddyext " } importName = strings.Replace(dir.Name, "-", "_", -1) begin = strings.Replace(begin, fmt.Sprintf(`"{{import-%s}}"`, importName), fmt.Sprintf(`%s%s "%s"`, comment, importName, dir.ImportPath), -1) if dir.Core == true { dirOrder = dirOrder + fmt.Sprintf(` %s{"%s", %s},`+"\n", comment, dir.Name, dir.Setup) } else { dirOrder = dirOrder + fmt.Sprintf(` %s{"%s", %s.Setup},`+"\n", comment, dir.Name, importName) } } out = begin + dirOrderStart + dirOrder + dirOrderEnd + end return ioutil.WriteFile(d.file, []byte(out), os.FileMode(0660)) }
// gen is code generation function that insert custom directives at runtime. func gen(middlewares features.Middlewares) custombuild.CodeGenFunc { return func(src string, packages []string) (err error) { // prevent possible panic from assertions. defer func() { if recover() != nil { err = errParse } }() // if no middleware is added, no code generation needed. if len(middlewares) == 0 { return nil } fset := token.NewFileSet() file := filepath.Join(src, directivesFile) f, err := parser.ParseFile(fset, file, nil, 0) if err != nil { return err } packageNames, err := getPackageNames(middlewares.Packages()) if err != nil { return err } for _, m := range middlewares { astutil.AddImport(fset, f, m.Package) } var buf bytes.Buffer err = printer.Fprint(&buf, fset, f) if err != nil { return err } out := buf.String() for _, mid := range middlewares { f, err = parser.ParseFile(token.NewFileSet(), "", out, 0) node, ok := f.Scope.Lookup("directiveOrder").Decl.(ast.Node) if !ok { return errParse } snippet := fmt.Sprintf(`{"%s", %s.Setup},`+"\n", mid.Directive, packageNames[mid.Package]) // add to end of directives. end := int(node.End()) - 2 // if after is set, locate directive and add after it. after := getPrevDirective(mid.Directive) if after != "" { found := false c := node.(*ast.ValueSpec).Values[0].(*ast.CompositeLit) for _, m := range c.Elts { directive := m.(*ast.CompositeLit).Elts[0].(*ast.BasicLit) if strconv.Quote(after) == directive.Value && !found { end = int(m.End()) + 1 found = true } // check if directive exists if strconv.Quote(mid.Directive) == directive.Value { return fmt.Errorf("Directive '%s' exists in Caddy core, use a distinct name.", mid.Directive) } } if !found { return fmt.Errorf("Cannot place afer %s, directive '%s' not found.", config.After, config.After) } } out = out[:end] + snippet + out[end:] } return ioutil.WriteFile(file, []byte(out), os.FileMode(0660)) } }