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>

diff --git a/src/go/printer/printer.go b/src/go/printer/printer.go
index 9d738f41b4741d6d456793502ca3a361b2c192d7..57f9716f48559da31ee60184e57367ff1453b855 100644 (file)
--- 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 409a53fd29acc88e241a87ace16ca4982c7183b7..5984d2c4d24e1911b0d802c1150417032087574a 100644 (file)
--- 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)
+       }
+}