protogen: move comment generation into protogen

Most plugins need to copy comments from .proto source files into the
generated code. Move this functionality into protogen to avoid
duplicating it everywhere.

Change-Id: I48a96ba794192e7ddc00281342afd4805ef6fe0f
Reviewed-on: https://go-review.googlesource.com/c/142890
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>
diff --git a/protogen/protogen.go b/protogen/protogen.go
index 22f663e..0a80aaa 100644
--- a/protogen/protogen.go
+++ b/protogen/protogen.go
@@ -13,6 +13,7 @@
 import (
 	"bufio"
 	"bytes"
+	"encoding/binary"
 	"fmt"
 	"go/ast"
 	"go/parser"
@@ -390,6 +391,8 @@
 	// For example, the source file "dir/foo.proto" might have a filename prefix
 	// of "dir/foo". Appending ".pb.go" produces an output file of "dir/foo.pb.go".
 	GeneratedFilenamePrefix string
+
+	sourceInfo map[pathKey][]*descpb.SourceCodeInfo_Location
 }
 
 func newFile(gen *Plugin, p *descpb.FileDescriptorProto, packageName GoPackageName, importPath GoImportPath) (*File, error) {
@@ -405,6 +408,7 @@
 		Proto:         p,
 		GoPackageName: packageName,
 		GoImportPath:  importPath,
+		sourceInfo:    make(map[pathKey][]*descpb.SourceCodeInfo_Location),
 	}
 
 	// Determine the prefix for generated Go files.
@@ -425,6 +429,10 @@
 	}
 	f.GeneratedFilenamePrefix = prefix
 
+	for _, loc := range p.GetSourceCodeInfo().GetLocation() {
+		key := newPathKey(loc.Path)
+		f.sourceInfo[key] = append(f.sourceInfo[key], loc)
+	}
 	for i, mdescs := 0, desc.Messages(); i < mdescs.Len(); i++ {
 		f.Messages = append(f.Messages, newMessage(gen, f, nil, mdescs.Get(i)))
 	}
@@ -854,6 +862,29 @@
 	fmt.Fprintln(&g.buf)
 }
 
+// PrintLeadingComments writes the comment appearing before a location in
+// the .proto source to the generated file.
+//
+// It returns true if a comment was present at the location.
+func (g *GeneratedFile) PrintLeadingComments(loc Location) (hasComment bool) {
+	f := g.gen.filesByName[loc.SourceFile]
+	if f == nil {
+		return false
+	}
+	for _, infoLoc := range f.sourceInfo[newPathKey(loc.Path)] {
+		if infoLoc.LeadingComments == nil {
+			continue
+		}
+		for _, line := range strings.Split(strings.TrimSuffix(infoLoc.GetLeadingComments(), "\n"), "\n") {
+			g.buf.WriteString("//")
+			g.buf.WriteString(line)
+			g.buf.WriteString("\n")
+		}
+		return true
+	}
+	return false
+}
+
 // QualifiedGoIdent returns the string to use for a Go identifier.
 //
 // If the identifier is from a different Go package than the generated file,
@@ -1070,3 +1101,17 @@
 		Path:       n,
 	}
 }
+
+// A pathKey is a representation of a location path suitable for use as a map key.
+type pathKey struct {
+	s string
+}
+
+// newPathKey converts a location path to a pathKey.
+func newPathKey(path []int32) pathKey {
+	buf := make([]byte, 4*len(path))
+	for i, x := range path {
+		binary.LittleEndian.PutUint32(buf[i*4:], uint32(x))
+	}
+	return pathKey{string(buf)}
+}