builders/cover: register coverage without changing line numbers

This copies in Go's cmd/internal/edit package for bytebuffer-based
editing.

Fixes #2990

go/tools/builders/BUILD.bazel

@@ -11,6 +11,18 @@ go_test(
     ],
 )
 
+go_test(
+    name = "cover_test",
+    size = "small",
+    srcs = [
+        "cover.go",
+        "cover_test.go",
+        "edit.go",
+        "env.go",
+        "flags.go",
+    ],
+)
+
 filegroup(
     name = "builder_srcs",
     srcs = [
@@ -21,14 +33,14 @@ filegroup(
         "compile.go",
         "compilepkg.go",
         "cover.go",
+        "edit.go",
         "embedcfg.go",
         "env.go",
         "filter.go",
         "filter_buildid.go",
         "flags.go",
         "generate_nogo_main.go",
         "generate_test_main.go",
-        "imports.go",
         "importcfg.go",
         "link.go",
         "pack.go",

go/tools/builders/cover.go

@@ -18,10 +18,10 @@ import (
 	"bytes"
 	"flag"
 	"fmt"
-	"go/format"
 	"go/parser"
 	"go/token"
 	"io/ioutil"
+	"os"
 	"strconv"
 )
 
@@ -70,39 +70,57 @@ func instrumentForCoverage(goenv *env, srcPath, srcName, coverVar, mode, outPath
 	if err := goenv.runCommand(goargs); err != nil {
 		return err
 	}
+	coverSrc, err := os.ReadFile(outPath)
+	if err != nil {
+		return fmt.Errorf("instrumentForCoverage: reading instrumented source: %w", err)
+	}
 
-	return registerCoverage(outPath, coverVar, srcName)
+	coverSrc, err = registerCoverage(outPath, coverSrc, coverVar, srcName)
+	if err := ioutil.WriteFile(outPath, coverSrc, 0666); err != nil {
+		return fmt.Errorf("instrumentForCoverage: %v", err)
+	}
+	return nil
 }
 
-// registerCoverage modifies coverSrc, the output file from go tool cover. It
-// adds a call to coverdata.RegisterCoverage, which ensures the coverage
+// registerCoverage returns a modified coverSrc, the output file from go tool cover.
+// It adds a call to coverdata.RegisterCoverage, which ensures the coverage
 // data from each file is reported. The name by which the file is registered
 // need not match its original name (it may use the importpath).
-func registerCoverage(coverSrc, varName, srcName string) error {
+func registerCoverage(coverSrcFilename string, coverSrc []byte, varName, srcName string) ([]byte, error) {
 	// Parse the file.
 	fset := token.NewFileSet()
-	f, err := parser.ParseFile(fset, coverSrc, nil, parser.ParseComments)
+	f, err := parser.ParseFile(fset, coverSrcFilename, coverSrc, parser.ParseComments)
 	if err != nil {
-		return nil // parse error: proceed and let the compiler fail
+		return coverSrc, nil // parse error: proceed and let the compiler fail
 	}
 
-	// Ensure coverdata is imported in the AST. Use an existing import if present
+	// Perform edits using a byte buffer instead of the AST, because
+	// we can not use go/format to write the AST back out without
+	// changing line numbers.
+	editor := NewBuffer(coverSrc)
+
+	// Ensure coverdata is imported. Use an existing import if present
 	// or add a new one.
 	const coverdataPath = "github.com/bazelbuild/rules_go/go/tools/coverdata"
 	var coverdataName string
 	for _, imp := range f.Imports {
 		path, err := strconv.Unquote(imp.Path.Value)
 		if err != nil {
-			return nil // parse error: proceed and let the compiler fail
+			return coverSrc, nil // parse error: proceed and let the compiler fail
 		}
 		if path == coverdataPath {
 			if imp.Name != nil {
 				// renaming import
 				if imp.Name.Name == "_" {
 					// Change blank import to named import
-					imp.Name.Name = "coverdata"
+					editor.Replace(
+						fset.Position(imp.Name.Pos()).Offset,
+						fset.Position(imp.Name.End()).Offset,
+						"coverdata")
+					coverdataName = "coverdata"
+				} else {
+					coverdataName = imp.Name.Name
 				}
-				coverdataName = imp.Name.Name
 			} else {
 				// default import
 				coverdataName = "coverdata"
@@ -113,25 +131,18 @@ func registerCoverage(coverSrc, varName, srcName string) error {
 	if coverdataName == "" {
 		// No existing import. Add a new one.
 		coverdataName = "coverdata"
-		addNamedImport(fset, f, coverdataName, coverdataPath)
-	}
-	var buf bytes.Buffer
-	if err := format.Node(&buf, fset, f); err != nil {
-		return fmt.Errorf("registerCoverage: could not reformat coverage source %s: %v", coverSrc, err)
+		editor.Insert(fset.Position(f.Name.End()).Offset, fmt.Sprintf("; import %q", coverdataPath))
 	}
 
 	// Append an init function.
-	fmt.Fprintf(&buf, `
+	var buf = bytes.NewBuffer(editor.Bytes())
+	fmt.Fprintf(buf, `
 func init() {
 	%s.RegisterFile(%q,
 		%[3]s.Count[:],
 		%[3]s.Pos[:],
 		%[3]s.NumStmt[:])
 }
 `, coverdataName, srcName, varName)
-	if err := ioutil.WriteFile(coverSrc, buf.Bytes(), 0666); err != nil {
-		return fmt.Errorf("registerCoverage: %v", err)
-	}
-
-	return nil
+	return buf.Bytes(), nil
 }

go/tools/builders/cover_test.go

@@ -0,0 +1,116 @@
+package main
+
+import "testing"
+
+type test struct {
+	name string
+	in   string
+	out  string
+}
+
+var tests = []test{
+	{
+		name: "no imports",
+		in: `package main
+`,
+		out: `package main; import "github.com/bazelbuild/rules_go/go/tools/coverdata"
+
+func init() {
+	coverdata.RegisterFile("srcName",
+		varName.Count[:],
+		varName.Pos[:],
+		varName.NumStmt[:])
+}
+`,
+	},
+	{
+		name: "other imports",
+		in: `package main
+
+import (
+	"os"
+)
+`,
+		out: `package main; import "github.com/bazelbuild/rules_go/go/tools/coverdata"
+
+import (
+	"os"
+)
+
+func init() {
+	coverdata.RegisterFile("srcName",
+		varName.Count[:],
+		varName.Pos[:],
+		varName.NumStmt[:])
+}
+`,
+	},
+	{
+		name: "existing import",
+		in: `package main
+
+import "github.com/bazelbuild/rules_go/go/tools/coverdata"
+`,
+		out: `package main
+
+import "github.com/bazelbuild/rules_go/go/tools/coverdata"
+
+func init() {
+	coverdata.RegisterFile("srcName",
+		varName.Count[:],
+		varName.Pos[:],
+		varName.NumStmt[:])
+}
+`,
+	},
+	{
+		name: "existing _ import",
+		in: `package main
+
+import _ "github.com/bazelbuild/rules_go/go/tools/coverdata"
+`,
+		out: `package main
+
+import coverdata "github.com/bazelbuild/rules_go/go/tools/coverdata"
+
+func init() {
+	coverdata.RegisterFile("srcName",
+		varName.Count[:],
+		varName.Pos[:],
+		varName.NumStmt[:])
+}
+`,
+	},
+	{
+		name: "existing renamed import",
+		in: `package main
+
+import cover0 "github.com/bazelbuild/rules_go/go/tools/coverdata"
+`,
+		out: `package main
+
+import cover0 "github.com/bazelbuild/rules_go/go/tools/coverdata"
+
+func init() {
+	cover0.RegisterFile("srcName",
+		varName.Count[:],
+		varName.Pos[:],
+		varName.NumStmt[:])
+}
+`,
+	},
+}
+
+func TestRegisterCoverage(t *testing.T) {
+	for _, test := range tests {
+		coverSrc, err := registerCoverage("in.go", []byte(test.in), "varName", "srcName")
+		if err != nil {
+			t.Errorf("%q: %+v", test.name, err)
+			continue
+		}
+
+		if got, want := string(coverSrc), test.out; got != want {
+			t.Errorf("%q: got %v, want %v", test.name, got, want)
+		}
+	}
+}