builders/cover: register coverage without changing line numbers

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"
 )
 
@@ -74,19 +74,29 @@ func instrumentForCoverage(goenv *env, srcPath, srcName, coverVar, mode, outPath
 	return registerCoverage(outPath, coverVar, srcName)
 }
 
-// registerCoverage modifies coverSrc, the output file from go tool cover. It
-// adds a call to coverdata.RegisterCoverage, which ensures the coverage
+// registerCoverage modifies coverSrcFilename, 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, varName, srcName string) error {
+	coverSrc, err := os.ReadFile(coverSrcFilename)
+	if err != nil {
+		return fmt.Errorf("instrumentForCoverage: reading instrumented source: %w", err)
+	}
+
 	// 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
 	}
 
-	// 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
@@ -100,9 +110,14 @@ func registerCoverage(coverSrc, varName, srcName string) error {
 				// 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 +128,21 @@ 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 {
+	if err := ioutil.WriteFile(coverSrcFilename, buf.Bytes(), 0666); err != nil {
 		return fmt.Errorf("registerCoverage: %v", err)
 	}
-
 	return nil
 }

go/tools/builders/cover_test.go

@@ -0,0 +1,130 @@
+package main
+
+import (
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"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) {
+	var filename = filepath.Join(t.TempDir(), "test_input.go")
+	for _, test := range tests {
+		if err := ioutil.WriteFile(filename, []byte(test.in), 0666); err != nil {
+			t.Errorf("writing input file: %v", err)
+			return
+		}
+		err := registerCoverage(filename, "varName", "srcName")
+		if err != nil {
+			t.Errorf("%q: %+v", test.name, err)
+			continue
+		}
+		coverSrc, err := os.ReadFile(filename)
+		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)
+		}
+	}
+}

go/tools/builders/edit.go

@@ -0,0 +1,95 @@
+// Copyright 2017 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Copied from go1.17 tree: //src/cmd/internal/edit/edit.go
+
+// Package edit implements buffered position-based editing of byte slices.
+package main
+
+import (
+	"fmt"
+	"sort"
+)
+
+// A Buffer is a queue of edits to apply to a given byte slice.
+type Buffer struct {
+	old []byte
+	q   edits
+}
+
+// An edit records a single text modification: change the bytes in [start,end) to new.
+type edit struct {
+	start int
+	end   int
+	new   string
+}
+
+// An edits is a list of edits that is sortable by start offset, breaking ties by end offset.
+type edits []edit
+
+func (x edits) Len() int      { return len(x) }
+func (x edits) Swap(i, j int) { x[i], x[j] = x[j], x[i] }
+func (x edits) Less(i, j int) bool {
+	if x[i].start != x[j].start {
+		return x[i].start < x[j].start
+	}
+	return x[i].end < x[j].end
+}
+
+// NewBuffer returns a new buffer to accumulate changes to an initial data slice.
+// The returned buffer maintains a reference to the data, so the caller must ensure
+// the data is not modified until after the Buffer is done being used.
+func NewBuffer(data []byte) *Buffer {
+	return &Buffer{old: data}
+}
+
+func (b *Buffer) Insert(pos int, new string) {
+	if pos < 0 || pos > len(b.old) {
+		panic("invalid edit position")
+	}
+	b.q = append(b.q, edit{pos, pos, new})
+}
+
+func (b *Buffer) Delete(start, end int) {
+	if end < start || start < 0 || end > len(b.old) {
+		panic("invalid edit position")
+	}
+	b.q = append(b.q, edit{start, end, ""})
+}
+
+func (b *Buffer) Replace(start, end int, new string) {
+	if end < start || start < 0 || end > len(b.old) {
+		panic("invalid edit position")
+	}
+	b.q = append(b.q, edit{start, end, new})
+}
+
+// Bytes returns a new byte slice containing the original data
+// with the queued edits applied.
+func (b *Buffer) Bytes() []byte {
+	// Sort edits by starting position and then by ending position.
+	// Breaking ties by ending position allows insertions at point x
+	// to be applied before a replacement of the text at [x, y).
+	sort.Stable(b.q)
+
+	var new []byte
+	offset := 0
+	for i, e := range b.q {
+		if e.start < offset {
+			e0 := b.q[i-1]
+			panic(fmt.Sprintf("overlapping edits: [%d,%d)->%q, [%d,%d)->%q", e0.start, e0.end, e0.new, e.start, e.end, e.new))
+		}
+		new = append(new, b.old[offset:e.start]...)
+		offset = e.end
+		new = append(new, e.new...)
+	}
+	new = append(new, b.old[offset:]...)
+	return new
+}
+
+// String returns a string containing the original data
+// with the queued edits applied.
+func (b *Buffer) String() string {
+	return string(b.Bytes())
+}