// 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, 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 }
func fixImports(fset *token.FileSet, f *ast.File) (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) // 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, `\"`)) 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) // 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) 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) } } // Nil out any unused ImportSpecs, to be removed in following passes unusedImport := map[string]bool{} for pkg, is := range decls { if refs[pkg] == nil && pkg != "_" && pkg != "." { unusedImport[strings.Trim(is.Path.Value, `"`)] = true } } for ipath := range unusedImport { if ipath == "C" { // Don't remove cgo stuff. continue } astutil.DeleteImport(fset, f, ipath) } return added, nil }