Example #1
0
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
}
Example #2
0
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
}