1*9bb1b549SSpandan Das// Copyright 2017 The Bazel Authors. All rights reserved. 2*9bb1b549SSpandan Das// 3*9bb1b549SSpandan Das// Licensed under the Apache License, Version 2.0 (the "License"); 4*9bb1b549SSpandan Das// you may not use this file except in compliance with the License. 5*9bb1b549SSpandan Das// You may obtain a copy of the License at 6*9bb1b549SSpandan Das// 7*9bb1b549SSpandan Das// http://www.apache.org/licenses/LICENSE-2.0 8*9bb1b549SSpandan Das// 9*9bb1b549SSpandan Das// Unless required by applicable law or agreed to in writing, software 10*9bb1b549SSpandan Das// distributed under the License is distributed on an "AS IS" BASIS, 11*9bb1b549SSpandan Das// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12*9bb1b549SSpandan Das// See the License for the specific language governing permissions and 13*9bb1b549SSpandan Das// limitations under the License. 14*9bb1b549SSpandan Das 15*9bb1b549SSpandan Daspackage main 16*9bb1b549SSpandan Das 17*9bb1b549SSpandan Dasimport ( 18*9bb1b549SSpandan Das "bytes" 19*9bb1b549SSpandan Das "fmt" 20*9bb1b549SSpandan Das "go/parser" 21*9bb1b549SSpandan Das "go/token" 22*9bb1b549SSpandan Das "io/ioutil" 23*9bb1b549SSpandan Das "os" 24*9bb1b549SSpandan Das "strconv" 25*9bb1b549SSpandan Das) 26*9bb1b549SSpandan Das 27*9bb1b549SSpandan Das// instrumentForCoverage runs "go tool cover" on a source file to produce 28*9bb1b549SSpandan Das// a coverage-instrumented version of the file. It also registers the file 29*9bb1b549SSpandan Das// with the coverdata package. 30*9bb1b549SSpandan Dasfunc instrumentForCoverage(goenv *env, srcPath, srcName, coverVar, mode, outPath string) error { 31*9bb1b549SSpandan Das goargs := goenv.goTool("cover", "-var", coverVar, "-mode", mode, "-o", outPath, srcPath) 32*9bb1b549SSpandan Das if err := goenv.runCommand(goargs); err != nil { 33*9bb1b549SSpandan Das return err 34*9bb1b549SSpandan Das } 35*9bb1b549SSpandan Das 36*9bb1b549SSpandan Das return registerCoverage(outPath, coverVar, srcName) 37*9bb1b549SSpandan Das} 38*9bb1b549SSpandan Das 39*9bb1b549SSpandan Das// registerCoverage modifies coverSrcFilename, the output file from go tool cover. 40*9bb1b549SSpandan Das// It adds a call to coverdata.RegisterCoverage, which ensures the coverage 41*9bb1b549SSpandan Das// data from each file is reported. The name by which the file is registered 42*9bb1b549SSpandan Das// need not match its original name (it may use the importpath). 43*9bb1b549SSpandan Dasfunc registerCoverage(coverSrcFilename, varName, srcName string) error { 44*9bb1b549SSpandan Das coverSrc, err := os.ReadFile(coverSrcFilename) 45*9bb1b549SSpandan Das if err != nil { 46*9bb1b549SSpandan Das return fmt.Errorf("instrumentForCoverage: reading instrumented source: %w", err) 47*9bb1b549SSpandan Das } 48*9bb1b549SSpandan Das 49*9bb1b549SSpandan Das // Parse the file. 50*9bb1b549SSpandan Das fset := token.NewFileSet() 51*9bb1b549SSpandan Das f, err := parser.ParseFile(fset, coverSrcFilename, coverSrc, parser.ParseComments) 52*9bb1b549SSpandan Das if err != nil { 53*9bb1b549SSpandan Das return nil // parse error: proceed and let the compiler fail 54*9bb1b549SSpandan Das } 55*9bb1b549SSpandan Das 56*9bb1b549SSpandan Das // Perform edits using a byte buffer instead of the AST, because 57*9bb1b549SSpandan Das // we can not use go/format to write the AST back out without 58*9bb1b549SSpandan Das // changing line numbers. 59*9bb1b549SSpandan Das editor := NewBuffer(coverSrc) 60*9bb1b549SSpandan Das 61*9bb1b549SSpandan Das // Ensure coverdata is imported. Use an existing import if present 62*9bb1b549SSpandan Das // or add a new one. 63*9bb1b549SSpandan Das const coverdataPath = "github.com/bazelbuild/rules_go/go/tools/coverdata" 64*9bb1b549SSpandan Das var coverdataName string 65*9bb1b549SSpandan Das for _, imp := range f.Imports { 66*9bb1b549SSpandan Das path, err := strconv.Unquote(imp.Path.Value) 67*9bb1b549SSpandan Das if err != nil { 68*9bb1b549SSpandan Das return nil // parse error: proceed and let the compiler fail 69*9bb1b549SSpandan Das } 70*9bb1b549SSpandan Das if path == coverdataPath { 71*9bb1b549SSpandan Das if imp.Name != nil { 72*9bb1b549SSpandan Das // renaming import 73*9bb1b549SSpandan Das if imp.Name.Name == "_" { 74*9bb1b549SSpandan Das // Change blank import to named import 75*9bb1b549SSpandan Das editor.Replace( 76*9bb1b549SSpandan Das fset.Position(imp.Name.Pos()).Offset, 77*9bb1b549SSpandan Das fset.Position(imp.Name.End()).Offset, 78*9bb1b549SSpandan Das "coverdata") 79*9bb1b549SSpandan Das coverdataName = "coverdata" 80*9bb1b549SSpandan Das } else { 81*9bb1b549SSpandan Das coverdataName = imp.Name.Name 82*9bb1b549SSpandan Das } 83*9bb1b549SSpandan Das } else { 84*9bb1b549SSpandan Das // default import 85*9bb1b549SSpandan Das coverdataName = "coverdata" 86*9bb1b549SSpandan Das } 87*9bb1b549SSpandan Das break 88*9bb1b549SSpandan Das } 89*9bb1b549SSpandan Das } 90*9bb1b549SSpandan Das if coverdataName == "" { 91*9bb1b549SSpandan Das // No existing import. Add a new one. 92*9bb1b549SSpandan Das coverdataName = "coverdata" 93*9bb1b549SSpandan Das editor.Insert(fset.Position(f.Name.End()).Offset, fmt.Sprintf("; import %q", coverdataPath)) 94*9bb1b549SSpandan Das } 95*9bb1b549SSpandan Das 96*9bb1b549SSpandan Das // Append an init function. 97*9bb1b549SSpandan Das var buf = bytes.NewBuffer(editor.Bytes()) 98*9bb1b549SSpandan Das fmt.Fprintf(buf, ` 99*9bb1b549SSpandan Dasfunc init() { 100*9bb1b549SSpandan Das %s.RegisterFile(%q, 101*9bb1b549SSpandan Das %[3]s.Count[:], 102*9bb1b549SSpandan Das %[3]s.Pos[:], 103*9bb1b549SSpandan Das %[3]s.NumStmt[:]) 104*9bb1b549SSpandan Das} 105*9bb1b549SSpandan Das`, coverdataName, srcName, varName) 106*9bb1b549SSpandan Das if err := ioutil.WriteFile(coverSrcFilename, buf.Bytes(), 0666); err != nil { 107*9bb1b549SSpandan Das return fmt.Errorf("registerCoverage: %v", err) 108*9bb1b549SSpandan Das } 109*9bb1b549SSpandan Das return nil 110*9bb1b549SSpandan Das} 111