compiler/protogen, cmd/protoc-gen-go: use alternative comments API

This is a breaking change. High-level protogen API changes:
* remove GeneratedFile.PrintLeadingComments method
* add {Message,Field,Oneof,Enum,EnumValue,Service,Method}.Comments field
* add CommentSet and Comments type

CL/183157 added protoreflect.SourceLocations and it was discovered
that there can actually be duplicate locations for certain paths.
For that reason, we decided not to expose any helper methods
for looking up locations by path since it is unclear which location
to return if multiple matches.

The protogen.GeneratedFile.PrintLeadingComments has a similar dilemma
where it also needs to figure out what to do when duplicates exist.
Previously, it just chooses the first one with comments,
which may not be the right choice in a given context.

Analysis of current PrintLeadingComments usage shows that it is only
ever used (except once) for descriptor declarations.
In the case of descriptor declarations, they are guaranteed by protoc
to have only location.

Thus, we avoid the duplicate location problem by:
* Providing a CommentSet for every descriptor. The CommentSet contains
a set of leading and trailing comments of the Comments type.
* The Comments.String method knows how to interpret the comments
as provided by protoc and format them as // prefixed line comments.
* Values of the Comments type can be passed to the P method.

We drop direct support printing leading comments for non-descriptor locations,
but the exposure of the Comments type makes it easy for users to manually
handle other types of comments themselves.

Change-Id: Id4851456dc4e64d76bd6a30e8ad6137408dfb27a
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/189198
Reviewed-by: Damien Neil <dneil@google.com>
diff --git a/cmd/protoc-gen-go-grpc/internal_gengogrpc/grpc.go b/cmd/protoc-gen-go-grpc/internal_gengogrpc/grpc.go
index 54048a3..fd25cd0 100644
--- a/cmd/protoc-gen-go-grpc/internal_gengogrpc/grpc.go
+++ b/cmd/protoc-gen-go-grpc/internal_gengogrpc/grpc.go
@@ -71,9 +71,9 @@
 	g.Annotate(clientName, service.Location)
 	g.P("type ", clientName, " interface {")
 	for _, method := range service.Methods {
-		g.PrintLeadingComments(method.Location)
 		g.Annotate(clientName+"."+method.GoName, method.Location)
-		g.P(clientSignature(g, method))
+		g.P(method.Comments.Leading,
+			clientSignature(g, method))
 	}
 	g.P("}")
 	g.P()
@@ -117,9 +117,9 @@
 	g.Annotate(serverType, service.Location)
 	g.P("type ", serverType, " interface {")
 	for _, method := range service.Methods {
-		g.PrintLeadingComments(method.Location)
 		g.Annotate(serverType+"."+method.GoName, method.Location)
-		g.P(serverSignature(g, method))
+		g.P(method.Comments.Leading,
+			serverSignature(g, method))
 	}
 	g.P("}")
 	g.P()
diff --git a/cmd/protoc-gen-go/internal_gengo/main.go b/cmd/protoc-gen-go/internal_gengo/main.go
index be2a240..eeffe73 100644
--- a/cmd/protoc-gen-go/internal_gengo/main.go
+++ b/cmd/protoc-gen-go/internal_gengo/main.go
@@ -151,11 +151,15 @@
 		g.P("// source: ", f.Desc.Path())
 	}
 	g.P()
-	g.PrintLeadingComments(protogen.Location{
-		SourceFile: f.Proto.GetName(),
-		Path:       []int32{fieldnum.FileDescriptorProto_Package},
-	})
-	g.P()
+
+	for _, loc := range f.Proto.GetSourceCodeInfo().GetLocation() {
+		if len(loc.Path) == 1 && loc.Path[0] == fieldnum.FileDescriptorProto_Package {
+			if s := loc.GetLeadingComments(); s != "" {
+				g.P(protogen.Comments(s))
+				g.P()
+			}
+		}
+	}
 	g.P("package ", f.GoPackageName)
 	g.P()
 
@@ -274,17 +278,17 @@
 
 func genEnum(gen *protogen.Plugin, g *protogen.GeneratedFile, f *fileInfo, enum *protogen.Enum) {
 	// Enum type declaration.
-	g.PrintLeadingComments(enum.Location)
 	g.Annotate(enum.GoIdent.GoName, enum.Location)
-	g.P("type ", enum.GoIdent, " int32",
+	g.P(enum.Comments.Leading,
+		"type ", enum.GoIdent, " int32",
 		deprecationComment(enum.Desc.Options().(*descriptorpb.EnumOptions).GetDeprecated()))
 
 	// Enum value constants.
 	g.P("const (")
 	for _, value := range enum.Values {
-		g.PrintLeadingComments(value.Location)
 		g.Annotate(value.GoIdent.GoName, value.Location)
-		g.P(value.GoIdent, " ", enum.GoIdent, " = ", value.Desc.Number(),
+		g.P(value.Comments.Leading,
+			value.GoIdent, " ", enum.GoIdent, " = ", value.Desc.Number(),
 			deprecationComment(value.Desc.Options().(*descriptorpb.EnumValueOptions).GetDeprecated()))
 	}
 	g.P(")")
@@ -377,15 +381,17 @@
 	}
 
 	// Message type declaration.
-	hasComment := g.PrintLeadingComments(message.Location)
+	leadingComments := message.Comments.Leading
 	if message.Desc.Options().(*descriptorpb.MessageOptions).GetDeprecated() {
-		if hasComment {
-			g.P("//")
+		if leadingComments != "" {
+			g.P(leadingComments, "//")
+			leadingComments = "" // avoid printing them again later
 		}
 		g.P(deprecationComment(true))
 	}
 	g.Annotate(message.GoIdent.GoName, message.Location)
-	g.P("type ", message.GoIdent, " struct {")
+	g.P(leadingComments,
+		"type ", message.GoIdent, " struct {")
 	genMessageFields(g, f, message)
 	g.P("}")
 	g.P()
@@ -446,20 +452,19 @@
 		if oneof.Fields[0] != field {
 			return // only generate for first appearance
 		}
-		if g.PrintLeadingComments(oneof.Location) {
-			g.P("//")
+		if oneof.Comments.Leading != "" {
+			g.P(oneof.Comments.Leading, "//")
 		}
 		g.P("// Types that are valid to be assigned to ", oneof.GoName, ":")
 		for _, field := range oneof.Fields {
-			g.PrintLeadingComments(field.Location)
-			g.P("//\t*", fieldOneofType(field))
+			g.P(field.Comments.Leading,
+				"//\t*", fieldOneofType(field))
 		}
 		g.Annotate(message.GoIdent.GoName+"."+oneof.GoName, oneof.Location)
 		g.P(oneof.GoName, " ", oneofInterfaceName(oneof), " `protobuf_oneof:\"", oneof.Desc.Name(), "\"`")
 		sf.append(oneof.GoName)
 		return
 	}
-	g.PrintLeadingComments(field.Location)
 	goType, pointer := fieldGoType(g, f, field)
 	if pointer {
 		goType = "*" + goType
@@ -482,7 +487,8 @@
 		name = "XXX_weak_" + name
 	}
 	g.Annotate(message.GoIdent.GoName+"."+name, field.Location)
-	g.P(name, " ", goType, " `", strings.Join(tags, " "), "`",
+	g.P(field.Comments.Leading,
+		name, " ", goType, " `", strings.Join(tags, " "), "`",
 		deprecationComment(field.Desc.Options().(*descriptorpb.FieldOptions).GetDeprecated()))
 	sf.append(field.GoName)
 }
diff --git a/compiler/protogen/protogen.go b/compiler/protogen/protogen.go
index 9a561f8..5d4e73c 100644
--- a/compiler/protogen/protogen.go
+++ b/compiler/protogen/protogen.go
@@ -399,7 +399,7 @@
 	// of "dir/foo". Appending ".pb.go" produces an output file of "dir/foo.pb.go".
 	GeneratedFilenamePrefix string
 
-	sourceInfo map[pathKey][]*descriptorpb.SourceCodeInfo_Location
+	comments map[pathKey]CommentSet
 }
 
 func newFile(gen *Plugin, p *descriptorpb.FileDescriptorProto, packageName GoPackageName, importPath GoImportPath) (*File, error) {
@@ -415,7 +415,7 @@
 		Proto:         p,
 		GoPackageName: packageName,
 		GoImportPath:  importPath,
-		sourceInfo:    make(map[pathKey][]*descriptorpb.SourceCodeInfo_Location),
+		comments:      make(map[pathKey]CommentSet),
 	}
 
 	// Determine the prefix for generated Go files.
@@ -441,8 +441,17 @@
 	f.GeneratedFilenamePrefix = prefix
 
 	for _, loc := range p.GetSourceCodeInfo().GetLocation() {
-		key := newPathKey(loc.Path)
-		f.sourceInfo[key] = append(f.sourceInfo[key], loc)
+		// Descriptors declarations are guaranteed to have unique comment sets.
+		// Other locations may not be unique, but we don't use them.
+		var leadingDetached []Comments
+		for _, s := range loc.GetLeadingDetachedComments() {
+			leadingDetached = append(leadingDetached, Comments(s))
+		}
+		f.comments[newPathKey(loc.Path)] = CommentSet{
+			LeadingDetached: leadingDetached,
+			Leading:         Comments(loc.GetLeadingComments()),
+			Trailing:        Comments(loc.GetTrailingComments()),
+		}
 	}
 	for i, mdescs := 0, desc.Messages(); i < mdescs.Len(); i++ {
 		f.Messages = append(f.Messages, newMessage(gen, f, nil, mdescs.Get(i)))
@@ -514,6 +523,7 @@
 	Enums      []*Enum      // nested enum declarations
 	Extensions []*Extension // nested extension declarations
 	Location   Location     // location of this message
+	Comments   CommentSet   // comments associated with this message
 }
 
 func newMessage(gen *Plugin, f *File, parent *Message, desc protoreflect.MessageDescriptor) *Message {
@@ -527,6 +537,7 @@
 		Desc:     desc,
 		GoIdent:  newGoIdent(f, desc),
 		Location: loc,
+		Comments: f.comments[newPathKey(loc.Path)],
 	}
 	gen.messagesByName[desc.FullName()] = message
 	for i, mdescs := 0, desc.Messages(); i < mdescs.Len(); i++ {
@@ -625,12 +636,13 @@
 	// '{{GoName}}' and a getter method named 'Get{{GoName}}'.
 	GoName string
 
-	Parent   *Message // message in which this field is defined; nil if top-level extension
-	Oneof    *Oneof   // containing oneof; nil if not part of a oneof
-	Extendee *Message // extended message for extension fields; nil otherwise
-	Enum     *Enum    // type for enum fields; nil otherwise
-	Message  *Message // type for message or group fields; nil otherwise
-	Location Location // location of this field
+	Parent   *Message   // message in which this field is defined; nil if top-level extension
+	Oneof    *Oneof     // containing oneof; nil if not part of a oneof
+	Extendee *Message   // extended message for extension fields; nil otherwise
+	Enum     *Enum      // type for enum fields; nil otherwise
+	Message  *Message   // type for message or group fields; nil otherwise
+	Location Location   // location of this field
+	Comments CommentSet // comments associated with this field
 }
 
 func newField(gen *Plugin, f *File, message *Message, desc protoreflect.FieldDescriptor) *Field {
@@ -648,6 +660,7 @@
 		GoName:   camelCase(string(desc.Name())),
 		Parent:   message,
 		Location: loc,
+		Comments: f.comments[newPathKey(loc.Path)],
 	}
 	if desc.ContainingOneof() != nil {
 		field.Oneof = message.Oneofs[desc.ContainingOneof().Index()]
@@ -695,15 +708,18 @@
 	Parent *Message // message in which this oneof occurs
 	Fields []*Field // fields that are part of this oneof
 
-	Location Location // location of this oneof
+	Location Location   // location of this oneof
+	Comments CommentSet // comments associated with this oneof
 }
 
 func newOneof(gen *Plugin, f *File, message *Message, desc protoreflect.OneofDescriptor) *Oneof {
+	loc := message.Location.appendPath(fieldnum.DescriptorProto_OneofDecl, int32(desc.Index()))
 	return &Oneof{
 		Desc:     desc,
 		Parent:   message,
 		GoName:   camelCase(string(desc.Name())),
-		Location: message.Location.appendPath(fieldnum.DescriptorProto_OneofDecl, int32(desc.Index())),
+		Location: loc,
+		Comments: f.comments[newPathKey(loc.Path)],
 	}
 }
 
@@ -720,7 +736,8 @@
 	GoIdent GoIdent      // name of the generated Go type
 	Values  []*EnumValue // enum values
 
-	Location Location // location of this enum
+	Location Location   // location of this enum
+	Comments CommentSet // comments associated with this enum
 }
 
 func newEnum(gen *Plugin, f *File, parent *Message, desc protoreflect.EnumDescriptor) *Enum {
@@ -734,6 +751,7 @@
 		Desc:     desc,
 		GoIdent:  newGoIdent(f, desc),
 		Location: loc,
+		Comments: f.comments[newPathKey(loc.Path)],
 	}
 	gen.enumsByName[desc.FullName()] = enum
 	for i, evdescs := 0, enum.Desc.Values(); i < evdescs.Len(); i++ {
@@ -748,7 +766,8 @@
 
 	GoIdent GoIdent // name of the generated Go type
 
-	Location Location // location of this enum value
+	Location Location   // location of this enum value
+	Comments CommentSet // comments associated with this enum value
 }
 
 func newEnumValue(gen *Plugin, f *File, message *Message, enum *Enum, desc protoreflect.EnumValueDescriptor) *EnumValue {
@@ -761,10 +780,12 @@
 		parentIdent = message.GoIdent
 	}
 	name := parentIdent.GoName + "_" + string(desc.Name())
+	loc := enum.Location.appendPath(fieldnum.EnumDescriptorProto_Value, int32(desc.Index()))
 	return &EnumValue{
 		Desc:     desc,
 		GoIdent:  f.GoImportPath.Ident(name),
-		Location: enum.Location.appendPath(fieldnum.EnumDescriptorProto_Value, int32(desc.Index())),
+		Location: loc,
+		Comments: f.comments[newPathKey(loc.Path)],
 	}
 }
 
@@ -775,14 +796,17 @@
 	GoName  string
 	Methods []*Method // service method definitions
 
-	Location Location // location of this service
+	Location Location   // location of this service
+	Comments CommentSet // comments associated with this service
 }
 
 func newService(gen *Plugin, f *File, desc protoreflect.ServiceDescriptor) *Service {
+	loc := f.location(fieldnum.FileDescriptorProto_Service, int32(desc.Index()))
 	service := &Service{
 		Desc:     desc,
 		GoName:   camelCase(string(desc.Name())),
-		Location: f.location(fieldnum.FileDescriptorProto_Service, int32(desc.Index())),
+		Location: loc,
+		Comments: f.comments[newPathKey(loc.Path)],
 	}
 	for i, mdescs := 0, desc.Methods(); i < mdescs.Len(); i++ {
 		service.Methods = append(service.Methods, newMethod(gen, f, service, mdescs.Get(i)))
@@ -799,15 +823,18 @@
 	Input  *Message
 	Output *Message
 
-	Location Location // location of this method
+	Location Location   // location of this method
+	Comments CommentSet // comments associated with this method
 }
 
 func newMethod(gen *Plugin, f *File, service *Service, desc protoreflect.MethodDescriptor) *Method {
+	loc := service.Location.appendPath(fieldnum.ServiceDescriptorProto_Method, int32(desc.Index()))
 	method := &Method{
 		Desc:     desc,
 		GoName:   camelCase(string(desc.Name())),
 		Parent:   service,
-		Location: service.Location.appendPath(fieldnum.ServiceDescriptorProto_Method, int32(desc.Index())),
+		Location: loc,
+		Comments: f.comments[newPathKey(loc.Path)],
 	}
 	return method
 }
@@ -882,29 +909,6 @@
 	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,
@@ -1161,3 +1165,30 @@
 	}
 	return pathKey{string(buf)}
 }
+
+// CommentSet is a set of leading and trailing comments associated
+// with a .proto descriptor declaration.
+type CommentSet struct {
+	LeadingDetached []Comments
+	Leading         Comments
+	Trailing        Comments
+}
+
+// Comments is a comments string as provided by protoc.
+type Comments string
+
+// String formats the comments by inserting // to the start of each line,
+// ensuring that there is a trailing newline.
+// An empty comment is formatted as an empty string.
+func (c Comments) String() string {
+	if c == "" {
+		return ""
+	}
+	var b []byte
+	for _, line := range strings.Split(strings.TrimSuffix(string(c), "\n"), "\n") {
+		b = append(b, "//"...)
+		b = append(b, line...)
+		b = append(b, "\n"...)
+	}
+	return string(b)
+}