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/cmd/protoc-gen-go/internal_gengo/main.go b/cmd/protoc-gen-go/internal_gengo/main.go
index c3154f5..509bbc0 100644
--- a/cmd/protoc-gen-go/internal_gengo/main.go
+++ b/cmd/protoc-gen-go/internal_gengo/main.go
@@ -32,7 +32,6 @@
 
 type fileInfo struct {
 	*protogen.File
-	locationMap   map[string][]*descpb.SourceCodeInfo_Location
 	descriptorVar string // var containing the gzipped FileDescriptorProto
 	allEnums      []*protogen.Enum
 	allMessages   []*protogen.Message
@@ -42,12 +41,7 @@
 // GenerateFile generates the contents of a .pb.go file.
 func GenerateFile(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile) {
 	f := &fileInfo{
-		File:        file,
-		locationMap: make(map[string][]*descpb.SourceCodeInfo_Location),
-	}
-	for _, loc := range file.Proto.GetSourceCodeInfo().GetLocation() {
-		key := pathKey(loc.Path)
-		f.locationMap[key] = append(f.locationMap[key], loc)
+		File: file,
 	}
 
 	// The different order for enums and extensions is to match the output
@@ -76,7 +70,10 @@
 	}
 	g.P()
 	const filePackageField = 2 // FileDescriptorProto.package
-	genComment(g, f, protogen.Location{Path: []int32{filePackageField}})
+	g.PrintLeadingComments(protogen.Location{
+		SourceFile: f.Proto.GetName(),
+		Path:       []int32{filePackageField},
+	})
 	g.P()
 	g.P("package ", f.GoPackageName)
 	g.P()
@@ -237,13 +234,13 @@
 }
 
 func genEnum(gen *protogen.Plugin, g *protogen.GeneratedFile, f *fileInfo, enum *protogen.Enum) {
-	genComment(g, f, enum.Location)
+	g.PrintLeadingComments(enum.Location)
 	g.Annotate(enum.GoIdent.GoName, enum.Location)
 	g.P("type ", enum.GoIdent, " int32",
 		deprecationComment(enumOptions(gen, enum).GetDeprecated()))
 	g.P("const (")
 	for _, value := range enum.Values {
-		genComment(g, f, value.Location)
+		g.PrintLeadingComments(value.Location)
 		g.Annotate(value.GoIdent.GoName, value.Location)
 		g.P(value.GoIdent, " ", enum.GoIdent, " = ", value.Desc.Number(),
 			deprecationComment(enumValueOptions(gen, value).GetDeprecated()))
@@ -335,7 +332,7 @@
 		return
 	}
 
-	hasComment := genComment(g, f, message.Location)
+	hasComment := g.PrintLeadingComments(message.Location)
 	if messageOptions(gen, message).GetDeprecated() {
 		if hasComment {
 			g.P("//")
@@ -355,7 +352,7 @@
 			}
 			continue
 		}
-		genComment(g, f, field.Location)
+		g.PrintLeadingComments(field.Location)
 		goType, pointer := fieldGoType(g, field)
 		if pointer {
 			goType = "*" + goType
@@ -902,20 +899,6 @@
 	}
 }
 
-func genComment(g *protogen.GeneratedFile, f *fileInfo, loc protogen.Location) (hasComment bool) {
-	for _, loc := range f.locationMap[pathKey(loc.Path)] {
-		if loc.LeadingComments == nil {
-			continue
-		}
-		for _, line := range strings.Split(strings.TrimSuffix(loc.GetLeadingComments(), "\n"), "\n") {
-			hasComment = true
-			g.P("//", line)
-		}
-		break
-	}
-	return hasComment
-}
-
 // deprecationComment returns a standard deprecation comment if deprecated is true.
 func deprecationComment(deprecated bool) string {
 	if !deprecated {
@@ -924,18 +907,6 @@
 	return "// Deprecated: Do not use."
 }
 
-// pathKey converts a location path to a string suitable for use as a map key.
-func pathKey(path []int32) string {
-	var buf []byte
-	for i, x := range path {
-		if i != 0 {
-			buf = append(buf, ',')
-		}
-		buf = strconv.AppendInt(buf, int64(x), 10)
-	}
-	return string(buf)
-}
-
 func genWellKnownType(g *protogen.GeneratedFile, ptr string, ident protogen.GoIdent, desc protoreflect.Descriptor) {
 	if wellKnownTypes[desc.FullName()] {
 		g.P("func (", ptr, ident, `) XXX_WellKnownType() string { return "`, desc.Name(), `" }`)