From 108746816ddf01ad0c2dbea08a1baef08bc47313 Mon Sep 17 00:00:00 2001
From: "David R. Jenni" <david.r.jenni@gmail.com>
Date: Tue, 23 Feb 2016 09:31:56 +0100
Subject: [PATCH] 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      | 12 +++++--
 go/ast/astutil/imports_test.go | 58 ++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/go/ast/astutil/imports.go b/go/ast/astutil/imports.go
index 7f9b1629..457cc6b5 100644
--- a/go/ast/astutil/imports.go
+++ b/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
diff --git a/go/ast/astutil/imports_test.go b/go/ast/astutil/imports_test.go
index 9134b193..0a1c7b65 100644
--- a/go/ast/astutil/imports_test.go
+++ b/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"
+)
 `,
 	},
 }