go/ast/astutil: keep comment position for existing imports.

Fixes golang/go#10337.

Change-Id: Ia21547bb57790c077ca01a2ad656086da3f2ba8a
Reviewed-on: https://go-review.googlesource.com/19822
Reviewed-by: Robert Griesemer <gri@golang.org>

go/ast/astutil/imports.go

@@ -117,9 +117,15 @@ func AddNamedImport(fset *token.FileSet, f *ast.File, name, ipath string) (added
 	impDecl.Specs[insertAt] = newImport
 	pos := impDecl.Pos()
 	if insertAt > 0 {
-		// Assign same position as the previous import,
-		// so that the sorter sees it as being in the same block.
-		pos = impDecl.Specs[insertAt-1].Pos()
+		// If there is a comment after an existing import, preserve the comment
+		// position by adding the new import after the comment.
+		if spec, ok := impDecl.Specs[insertAt-1].(*ast.ImportSpec); ok && spec.Comment != nil {
+			pos = spec.Comment.End()
+		} else {
+			// Assign same position as the previous import,
+			// so that the sorter sees it as being in the same block.
+			pos = impDecl.Specs[insertAt-1].Pos()
+		}
 	}
 	if newImport.Name != nil {
 		newImport.Name.NamePos = pos

go/ast/astutil/imports_test.go

@@ -338,6 +338,64 @@ import (
 
 	"rsc.io/p"
 )
+`,
+	},
+	// Issue 10337: Preserve comment position
+	{
+		name: "issue 10337",
+		pkg:  "fmt",
+		in: `package main
+
+import (
+	"bytes" // a
+	"log" // c
+)
+`,
+		out: `package main
+
+import (
+	"bytes" // a
+	"fmt"
+	"log" // c
+)
+`,
+	},
+	{
+		name: "issue 10337 new import at the start",
+		pkg:  "bytes",
+		in: `package main
+
+import (
+	"fmt" // b
+	"log" // c
+)
+`,
+		out: `package main
+
+import (
+	"bytes"
+	"fmt" // b
+	"log" // c
+)
+`,
+	},
+	{
+		name: "issue 10337 new import at the end",
+		pkg:  "log",
+		in: `package main
+
+import (
+	"bytes" // a
+	"fmt" // b
+)
+`,
+		out: `package main
+
+import (
+	"bytes" // a
+	"fmt"   // b
+	"log"
+)
 `,
 	},
 }