cmd/bundle: include non-associated comments

https://golang.org/cl/45117 doesn't handle comments which don't
belong to any declarations. This CL addresses it.

Fixes golang/go#20627

Change-Id: I81c7cdc070efc6cb9e9f38101d7b2b8909916ba6
Reviewed-on: https://go-review.googlesource.com/45190
Reviewed-by: Alan Donovan <adonovan@google.com>

cmd/bundle/main.go

@@ -349,6 +349,17 @@ func bundle(src, dst, dstpkg, prefix string) ([]byte, error) {
 			return true
 		})
 
+		last := f.Package
+		if len(f.Imports) > 0 {
+			imp := f.Imports[len(f.Imports)-1]
+			last = imp.End()
+			if imp.Comment != nil {
+				if e := imp.Comment.End(); e > last {
+					last = e
+				}
+			}
+		}
+
 		// Pretty-print package-level declarations.
 		// but no package or import declarations.
 		var buf bytes.Buffer
@@ -356,13 +367,23 @@ func bundle(src, dst, dstpkg, prefix string) ([]byte, error) {
 			if decl, ok := decl.(*ast.GenDecl); ok && decl.Tok == token.IMPORT {
 				continue
 			}
+
+			beg, end := sourceRange(decl)
+
+			printComments(&out, f.Comments, last, beg)
+
 			buf.Reset()
 			format.Node(&buf, lprog.Fset, &printer.CommentedNode{Node: decl, Comments: f.Comments})
 			// Remove each "@@@." in the output.
 			// TODO(adonovan): not hygienic.
 			out.Write(bytes.Replace(buf.Bytes(), []byte("@@@."), nil, -1))
+
+			last = printSameLineComment(&out, f.Comments, lprog.Fset, end)
+
 			out.WriteString("\n\n")
 		}
+
+		printLastComments(&out, f.Comments, last)
 	}
 
 	// Now format the entire thing.
@@ -374,6 +395,69 @@ func bundle(src, dst, dstpkg, prefix string) ([]byte, error) {
 	return result, nil
 }
 
+// sourceRange returns the [beg, end) interval of source code
+// belonging to decl (incl. associated comments).
+func sourceRange(decl ast.Decl) (beg, end token.Pos) {
+	beg = decl.Pos()
+	end = decl.End()
+
+	var doc, com *ast.CommentGroup
+
+	switch d := decl.(type) {
+	case *ast.GenDecl:
+		doc = d.Doc
+		if len(d.Specs) > 0 {
+			switch spec := d.Specs[len(d.Specs)-1].(type) {
+			case *ast.ValueSpec:
+				com = spec.Comment
+			case *ast.TypeSpec:
+				com = spec.Comment
+			}
+		}
+	case *ast.FuncDecl:
+		doc = d.Doc
+	}
+
+	if doc != nil {
+		beg = doc.Pos()
+	}
+	if com != nil && com.End() > end {
+		end = com.End()
+	}
+
+	return beg, end
+}
+
+func printComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos, end token.Pos) {
+	for _, cg := range comments {
+		if pos <= cg.Pos() && cg.Pos() < end {
+			for _, c := range cg.List {
+				fmt.Fprintln(out, c.Text)
+			}
+			fmt.Fprintln(out)
+		}
+	}
+}
+
+const infinity = 1 << 30
+
+func printLastComments(out *bytes.Buffer, comments []*ast.CommentGroup, pos token.Pos) {
+	printComments(out, comments, pos, infinity)
+}
+
+func printSameLineComment(out *bytes.Buffer, comments []*ast.CommentGroup, fset *token.FileSet, pos token.Pos) token.Pos {
+	tf := fset.File(pos)
+	for _, cg := range comments {
+		if pos <= cg.Pos() && tf.Line(cg.Pos()) == tf.Line(pos) {
+			for _, c := range cg.List {
+				fmt.Fprintln(out, c.Text)
+			}
+			return cg.End()
+		}
+	}
+	return pos
+}
+
 type flagFunc func(string)
 
 func (f flagFunc) Set(s string) error {

cmd/bundle/main_test.go

@@ -2,6 +2,8 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+// +build go1.9
+
 package main
 
 import (

cmd/bundle/testdata/out.golden

@@ -23,21 +23,38 @@ func init() { prefixfoo() }
 type prefixS struct {
 	prefixt
 	u int
-}
+} /* multi-line
+comment
+*/
+
+// non-associated comment
+
+/*
+	non-associated comment2
+*/
 
 // Function bar.
 func prefixbar(s *prefixS) {
 	fmt.Println(s.prefixt, s.u) // comment inside function
 }
 
-type prefixt int
+// file-end comment
 
-const prefixc = 1
+type prefixt int // type1
+
+// const1
+const prefixc = 1 // const2
 
 func prefixfoo() {
 	fmt.Println(importdecl.F())
 }
 
+// zinit
+const (
+	prefixz1 = iota // z1
+	prefixz2        // z2
+) // zend
+
 func prefixbaz() {
 	renamedfmt.Println()
 	renamedfmt2.Println()

cmd/bundle/testdata/src/initial/a.go

@@ -1,6 +1,6 @@
 package initial
 
-import "fmt"
+import "fmt" // this comment should not be visible
 
 // init functions are not renamed
 func init() { foo() }
@@ -9,9 +9,19 @@ func init() { foo() }
 type S struct {
 	t
 	u int
-}
+} /* multi-line
+comment
+*/
+
+// non-associated comment
+
+/*
+	non-associated comment2
+*/
 
 // Function bar.
 func bar(s *S) {
 	fmt.Println(s.t, s.u) // comment inside function
 }
+
+// file-end comment

cmd/bundle/testdata/src/initial/b.go

@@ -7,10 +7,17 @@ import (
 	"domain.name/importdecl"
 )
 
-type t int
+type t int // type1
 
-const c = 1
+// const1
+const c = 1 // const2
 
 func foo() {
 	fmt.Println(importdecl.F())
 }
+
+// zinit
+const (
+	z1 = iota // z1
+	z2        // z2
+) // zend