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>

imports/fix_test.go

@@ -864,6 +864,29 @@ func main() {
 	_ = p.P
 	_ = time.Parse
 }
+`,
+	},
+
+	{
+		name: "issue #12097",
+		in: `// a
+// b
+// c
+
+func main() {
+    _ = fmt.Println
+}`,
+		out: `package main
+
+import "fmt"
+
+// a
+// b
+// c
+
+func main() {
+	_ = fmt.Println
+}
 `,
 	},
 }

imports/imports.go

@@ -63,7 +63,6 @@ func Process(filename string, src []byte, opt *Options) ([]byte, error) {
 
 	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 @@ func parse(fset *token.FileSet, filename string, src []byte, opt *Options) (*ast
 
 	// 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 @@ func parse(fset *token.FileSet, filename string, src []byte, opt *Options) (*ast
 
 		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