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) + } +}