go/printer: handle associated comments for CommentedNode

Current CommentedNode cannot handle associated comments which satisfy
    node.End() < comment.Pos()

This CL solves it.

Fixes #20635

Change-Id: I58e2e3703999bb38a6ce37112e986c4b1b2eace0
Reviewed-on: https://go-review.googlesource.com/45292
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Robert Griesemer <gri@golang.org>

2 files changed, 78 insertions, 0 deletions

diff --git a/src/go/printer/printer.go b/src/go/printer/printer.go
index 9d738f41b4..57f9716f48 100644
--- a/src/go/printer/printer.go
+++ b/src/go/printer/printer.go
@@ -1043,6 +1043,28 @@ func getDoc(n ast.Node) *ast.CommentGroup {
 	return nil
 }
 
+func getLastComment(n ast.Node) *ast.CommentGroup {
+	switch n := n.(type) {
+	case *ast.Field:
+		return n.Comment
+	case *ast.ImportSpec:
+		return n.Comment
+	case *ast.ValueSpec:
+		return n.Comment
+	case *ast.TypeSpec:
+		return n.Comment
+	case *ast.GenDecl:
+		if len(n.Specs) > 0 {
+			return getLastComment(n.Specs[len(n.Specs)-1])
+		}
+	case *ast.File:
+		if len(n.Comments) > 0 {
+			return n.Comments[len(n.Comments)-1]
+		}
+	}
+	return nil
+}
+
 func (p *printer) printNode(node interface{}) error {
 	// unpack *CommentedNode, if any
 	var comments []*ast.CommentGroup
@@ -1066,6 +1088,11 @@ func (p *printer) printNode(node interface{}) error {
 		if doc := getDoc(n); doc != nil {
 			beg = doc.Pos()
 		}
+		if com := getLastComment(n); com != nil {
+			if e := com.End(); e > end {
+				end = e
+			}
+		}
 		// token.Pos values are global offsets, we can
 		// compare them directly
 		i := 0
diff --git a/src/go/printer/printer_test.go b/src/go/printer/printer_test.go
index 409a53fd29..5984d2c4d2 100644
--- a/src/go/printer/printer_test.go
+++ b/src/go/printer/printer_test.go
@@ -659,3 +659,54 @@ func _() {}
 		t.Error(err)
 	}
 }
+
+func TestCommentedNode(t *testing.T) {
+	const (
+		input = `package main
+
+func foo() {
+	// comment inside func
+}
+
+// leading comment
+type bar int // comment2
+
+`
+
+		foo = `func foo() {
+	// comment inside func
+}`
+
+		bar = `// leading comment
+type bar int	// comment2
+`
+	)
+
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, "input.go", input, parser.ParseComments)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var buf bytes.Buffer
+
+	err = Fprint(&buf, fset, &CommentedNode{Node: f.Decls[0], Comments: f.Comments})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if buf.String() != foo {
+		t.Errorf("got %q, want %q", buf.String(), foo)
+	}
+
+	buf.Reset()
+
+	err = Fprint(&buf, fset, &CommentedNode{Node: f.Decls[1], Comments: f.Comments})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if buf.String() != bar {
+		t.Errorf("got %q, want %q", buf.String(), bar)
+	}
+}