From 6c27c68f27f0d7b30dc15cd675853cc3d95be1e5 Mon Sep 17 00:00:00 2001
From: Antonio Antelo <aantelov87@gmail.com>
Date: Sat, 10 Feb 2018 12:48:53 +0100
Subject: [PATCH] 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 | 23 +++++++++++++++++++++++
 imports/imports.go  | 17 +++++++++++------
 2 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/imports/fix_test.go b/imports/fix_test.go
index bc18772cf0..83ceb02501 100644
--- a/imports/fix_test.go
+++ b/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
+}
 `,
 	},
 }
diff --git a/imports/imports.go b/imports/imports.go
index bd67f91df9..0623248744 100644
--- a/imports/imports.go
+++ b/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