func (in *instrumenter) instrumentFile(f *ast.File, fset *token.FileSet, pkgpath string) error { pkgObj := in.instrumented[pkgpath] pkgCreated := false if pkgObj == nil { pkgCreated = true pkgObj = gocov.RegisterPackage(f.Name.Name) // FIXME(axw) use full package path in.instrumented[pkgpath] = pkgObj } state := &state{fset, f, pkgObj, nil} ast.Walk(&funcVisitor{state}, f) // Count the number of import GenDecl's. They're always first. nImportDecls := 0 for _, decl := range f.Decls { if decl, ok := decl.(*ast.GenDecl); !ok || decl.Tok != token.IMPORT { break } nImportDecls++ } // Redirect imports of instrumented packages. in.redirectImports(f) // Insert variable declarations for registered objects. var vardecls []ast.Decl pkgvarname := fmt.Sprint(pkgObj) if pkgCreated { // FIXME(axw) use full package path value := makeCall("gocov.RegisterPackage", makeLit(f.Name.Name)) vardecls = append(vardecls, makeVarDecl(pkgvarname, value)) } for _, fn := range state.functions { fnvarname := fmt.Sprint(fn) value := makeCall(pkgvarname+".RegisterFunction", makeLit(fn.Name), makeLit(fn.File), makeLit(fn.Start), makeLit(fn.End)) vardecls = append(vardecls, makeVarDecl(fnvarname, value)) for _, stmt := range fn.Statements { varname := fmt.Sprint(stmt) value := makeCall( fnvarname+".RegisterStatement", makeLit(stmt.Start), makeLit(stmt.End)) vardecls = append(vardecls, makeVarDecl(varname, value)) } } if len(f.Decls) > 0 { vardecls = append(vardecls, f.Decls[nImportDecls:]...) f.Decls = append(f.Decls[:nImportDecls], vardecls...) } else { f.Decls = vardecls } // Add a "gocov" import. if pkgCreated { gocovImportSpec := &ast.ImportSpec{ Path: makeLit(gocovPackagePath).(*ast.BasicLit)} gocovImportGenDecl := &ast.GenDecl{ Tok: token.IMPORT, Specs: []ast.Spec{gocovImportSpec}} tail := make([]ast.Decl, len(f.Decls)-nImportDecls) copy(tail, f.Decls[nImportDecls:]) head := append(f.Decls[:nImportDecls], gocovImportGenDecl) f.Decls = append(head, tail...) } // Clear out all cached comments. This forces the AST printer to use // node comments instead, repositioning them correctly. f.Comments = nil return nil }
func (in *instrumenter) instrumentFile(filename string, f *ast.File, fset *token.FileSet, pkgpath string) error { pkgObj := in.instrumented[pkgpath] pkgCreated := false if pkgObj == nil { pkgCreated = true pkgObj = gocov.RegisterPackage(f.Name.Name) // FIXME(axw) use full package path in.instrumented[pkgpath] = pkgObj } state := &state{fset, f, pkgObj, nil} ast.Walk(&funcVisitor{state}, f) // Count the number of import GenDecl's. They're always first. nImportDecls := 0 for _, decl := range f.Decls { if decl, ok := decl.(*ast.GenDecl); !ok || decl.Tok != token.IMPORT { break } nImportDecls++ } // Redirect imports of instrumented packages. in.redirectImports(f) // Insert variable declarations for registered objects. var vardecls []ast.Decl pkgvarname := fmt.Sprint(pkgObj) if pkgCreated { // FIXME(axw) use full package path value := makeCall("gocov.RegisterPackage", makeLit(f.Name.Name)) vardecls = append(vardecls, makeVarDecl(pkgvarname, value)) } for _, fn := range state.functions { fnvarname := fmt.Sprint(fn) value := makeCall(pkgvarname+".RegisterFunction", makeLit(fn.Name), makeLit(fn.File), makeLit(fn.Start), makeLit(fn.End)) vardecls = append(vardecls, makeVarDecl(fnvarname, value)) for _, stmt := range fn.Statements { varname := fmt.Sprint(stmt) value := makeCall( fnvarname+".RegisterStatement", makeLit(stmt.Start), makeLit(stmt.End)) vardecls = append(vardecls, makeVarDecl(varname, value)) } } if len(f.Decls) > 0 { vardecls = append(vardecls, f.Decls[nImportDecls:]...) f.Decls = append(f.Decls[:nImportDecls], vardecls...) } else { f.Decls = vardecls } // Add a "gocov" import. if pkgCreated { gocovImportSpec := &ast.ImportSpec{ Path: makeLit(gocovPackagePath).(*ast.BasicLit)} gocovImportGenDecl := &ast.GenDecl{ Tok: token.IMPORT, Specs: []ast.Spec{gocovImportSpec}} tail := make([]ast.Decl, len(f.Decls)-nImportDecls) copy(tail, f.Decls[nImportDecls:]) head := append(f.Decls[:nImportDecls], gocovImportGenDecl) f.Decls = append(head, tail...) } // Record function comment associations var funcComments = make(map[string]*ast.FuncDecl, 0) for _, decl := range f.Decls { if n, ok := decl.(*ast.FuncDecl); ok && n.Doc != nil { funcComments[n.Name.String()] = n } } // Clear out all comments except for the comments attached to // existing import specs. if nImportDecls > 0 { end := f.Decls[nImportDecls-1].Pos() comments := make([]*ast.CommentGroup, 0, len(f.Comments)) for _, group := range f.Comments { if group.End() < end { comments = append(comments, group) } else { break } } f.Comments = comments } else { f.Comments = []*ast.CommentGroup{} } // Print and reparse to reinsert comments o := bytes.NewBuffer(make([]byte, 0, 512)) printer.Fprint(o, fset, f) parserMode := parser.DeclarationErrors | parser.ParseComments reparsed, err := parser.ParseFile(fset, filename, o, parserMode) if err != nil { return errors.New(fmt.Sprint("second pass parse error:", err)) } ric := &reinsertComments{funcs: funcComments, fileSet: fset, content: o.Bytes()} ast.Walk(ric, reparsed) reparsed, err = parser.ParseFile(fset, filename, ric.content, parserMode) *f = *reparsed return err }