imports: fix mangled comments after package clause insertion

Fixes golang/go#12097

Change-Id: Ie6a6aa997e89700e49d703b7fd00f515b03ad6f8
Reviewed-on: https://go-review.googlesource.com/93235
Reviewed-by: Josh Bleecher Snyder <josharian@gmail.com>
Run-TryBot: Josh Bleecher Snyder <josharian@gmail.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
diff --git a/imports/fix_test.go b/imports/fix_test.go
index bc18772..83ceb02 100644
--- a/imports/fix_test.go
+++ b/imports/fix_test.go
@@ -866,6 +866,29 @@
 }
 `,
 	},
+
+	{
+		name: "issue #12097",
+		in: `// a
+// b
+// c
+
+func main() {
+    _ = fmt.Println
+}`,
+		out: `package main
+
+import "fmt"
+
+// a
+// b
+// c
+
+func main() {
+	_ = fmt.Println
+}
+`,
+	},
 }
 
 func TestFixImports(t *testing.T) {
diff --git a/imports/imports.go b/imports/imports.go
index bd67f91..0623248 100644
--- a/imports/imports.go
+++ b/imports/imports.go
@@ -63,7 +63,6 @@
 
 	sortImports(fileSet, file)
 	imps := astutil.Imports(fileSet, file)
-
 	var spacesBefore []string // import paths we need spaces before
 	for _, impSection := range imps {
 		// Within each block of contiguous imports, see if any
@@ -136,11 +135,18 @@
 
 	// If this is a declaration list, make it a source file
 	// by inserting a package clause.
-	// Insert using a ;, not a newline, so that the line numbers
-	// in psrc match the ones in src.
-	psrc := append([]byte("package main;"), src...)
+	// Insert using a ;, not a newline, so that parse errors are on
+	// the correct line.
+	const prefix = "package main;"
+	psrc := append([]byte(prefix), src...)
 	file, err = parser.ParseFile(fset, filename, psrc, parserMode)
 	if err == nil {
+		// Gofmt will turn the ; into a \n.
+		// Do that ourselves now and update the file contents,
+		// so that positions and line numbers are correct going forward.
+		psrc[len(prefix)-1] = '\n'
+		fset.File(file.Package).SetLinesForContent(psrc)
+
 		// If a main function exists, we will assume this is a main
 		// package and leave the file.
 		if containsMainFunc(file) {
@@ -149,8 +155,7 @@
 
 		adjust := func(orig, src []byte) []byte {
 			// Remove the package clause.
-			// Gofmt has turned the ; into a \n.
-			src = src[len("package main\n"):]
+			src = src[len(prefix):]
 			return matchSpace(orig, src)
 		}
 		return file, adjust, nil