goprotobuf: Generate JSON marshal/unmarshal methods for MessageSet.

This only implements marshaling. Unmarshaling is much fiddlier,
and it's not even clear that people care about that, so I'm punting
it to a later date.

LGTM=gmlewis
R=gmlewis
CC=golang-codereviews
https://codereview.appspot.com/156950045

4 files changed, 94 insertions, 3 deletions

diff --git a/proto/message_set.go b/proto/message_set.go
index 1a17809..5504855 100644
--- a/proto/message_set.go
+++ b/proto/message_set.go
@@ -36,7 +36,10 @@ package proto
  */
 
 import (
+	"bytes"
+	"encoding/json"
 	"errors"
+	"fmt"
 	"reflect"
 	"sort"
 )
@@ -211,6 +214,61 @@ func UnmarshalMessageSet(buf []byte, m map[int32]Extension) error {
 	return nil
 }
 
+// MarshalMessageSetJSON encodes the extension map represented by m in JSON format.
+// It is called by generated MarshalJSON methods on protocol buffer messages with the message_set_wire_format option.
+func MarshalMessageSetJSON(m map[int32]Extension) ([]byte, error) {
+	var b bytes.Buffer
+	b.WriteByte('{')
+
+	// Process the map in key order for deterministic output.
+	ids := make([]int32, 0, len(m))
+	for id := range m {
+		ids = append(ids, id)
+	}
+	sort.Sort(int32Slice(ids)) // int32Slice defined in text.go
+
+	for i, id := range ids {
+		ext := m[id]
+		if i > 0 {
+			b.WriteByte(',')
+		}
+
+		msd, ok := messageSetMap[id]
+		if !ok {
+			// Unknown type; we can't render it, so skip it.
+			continue
+		}
+		fmt.Fprintf(&b, `"[%s]":`, msd.name)
+
+		x := ext.value
+		if x == nil {
+			x = reflect.New(msd.t.Elem()).Interface()
+			if err := Unmarshal(ext.enc, x.(Message)); err != nil {
+				return nil, err
+			}
+		}
+		d, err := json.Marshal(x)
+		if err != nil {
+			return nil, err
+		}
+		b.Write(d)
+	}
+	b.WriteByte('}')
+	return b.Bytes(), nil
+}
+
+// UnmarshalMessageSetJSON decodes the extension map encoded in buf in JSON format.
+// It is called by generated UnmarshalJSON methods on protocol buffer messages with the message_set_wire_format option.
+func UnmarshalMessageSetJSON(buf []byte, m map[int32]Extension) error {
+	// Common-case fast path.
+	if len(buf) == 0 || bytes.Equal(buf, []byte("{}")) {
+		return nil
+	}
+
+	// This is fairly tricky, and it's not clear that it is needed.
+	return errors.New("TODO: UnmarshalMessageSetJSON not yet implemented")
+}
+
 // A global registry of types that can be used in a MessageSet.
 
 var messageSetMap = make(map[int32]messageSetDesc)
@@ -221,9 +279,9 @@ type messageSetDesc struct {
 }
 
 // RegisterMessageSetType is called from the generated code.
-func RegisterMessageSetType(i messageTypeIder, name string) {
-	messageSetMap[i.MessageTypeId()] = messageSetDesc{
-		t:    reflect.TypeOf(i),
+func RegisterMessageSetType(m Message, fieldNum int32, name string) {
+	messageSetMap[fieldNum] = messageSetDesc{
+		t:    reflect.TypeOf(m),
 		name: name,
 	}
 }
diff --git a/protoc-gen-go/generator/generator.go b/protoc-gen-go/generator/generator.go
index a86a121..3912422 100644
--- a/protoc-gen-go/generator/generator.go
+++ b/protoc-gen-go/generator/generator.go
@@ -1523,6 +1523,16 @@ func (g *Generator) generateMessage(message *Descriptor) {
 			g.P("return ", g.Pkg["proto"], ".UnmarshalMessageSet(buf, m.ExtensionMap())")
 			g.Out()
 			g.P("}")
+			g.P("func (m *", ccTypeName, ") MarshalJSON() ([]byte, error) {")
+			g.In()
+			g.P("return ", g.Pkg["proto"], ".MarshalMessageSetJSON(m.XXX_extensions)")
+			g.Out()
+			g.P("}")
+			g.P("func (m *", ccTypeName, ") UnmarshalJSON(buf []byte) error {")
+			g.In()
+			g.P("return ", g.Pkg["proto"], ".UnmarshalMessageSetJSON(buf, m.XXX_extensions)")
+			g.Out()
+			g.P("}")
 			g.P("// ensure ", ccTypeName, " satisfies proto.Marshaler and proto.Unmarshaler")
 			g.P("var _ ", g.Pkg["proto"], ".Marshaler = (*", ccTypeName, ")(nil)")
 			g.P("var _ ", g.Pkg["proto"], ".Unmarshaler = (*", ccTypeName, ")(nil)")
@@ -1757,8 +1767,10 @@ func (g *Generator) generateExtension(ext *ExtensionDescriptor) {
 	// Special case for proto2 message sets: If this extension is extending
 	// proto2_bridge.MessageSet, and its final name component is "message_set_extension",
 	// then drop that last component.
+	mset := false
 	if extendedType == "*proto2_bridge.MessageSet" && typeName[len(typeName)-1] == "message_set_extension" {
 		typeName = typeName[:len(typeName)-1]
+		mset = true
 	}
 
 	// For text formatting, the package must be exactly what the .proto file declares,
@@ -1780,6 +1792,15 @@ func (g *Generator) generateExtension(ext *ExtensionDescriptor) {
 	g.P("}")
 	g.P()
 
+	if mset {
+		// Generate a bit more code to register with message_set.go.
+		g.P("func init() { ")
+		g.In()
+		g.P(g.Pkg["proto"], ".RegisterMessageSetType((", fieldType, ")(nil), ", field.Number, ", \"", extName, "\")")
+		g.Out()
+		g.P("}")
+	}
+
 	g.file.addExport(ext, constOrVarSymbol{ccTypeName, "var", ""})
 }
 
diff --git a/protoc-gen-go/testdata/my_test/test.pb.go b/protoc-gen-go/testdata/my_test/test.pb.go
index 1165dea..361a813 100644
--- a/protoc-gen-go/testdata/my_test/test.pb.go
+++ b/protoc-gen-go/testdata/my_test/test.pb.go
@@ -413,6 +413,12 @@ func (m *OldReply) Marshal() ([]byte, error) {
 func (m *OldReply) Unmarshal(buf []byte) error {
 	return proto.UnmarshalMessageSet(buf, m.ExtensionMap())
 }
+func (m *OldReply) MarshalJSON() ([]byte, error) {
+	return proto.MarshalMessageSetJSON(m.XXX_extensions)
+}
+func (m *OldReply) UnmarshalJSON(buf []byte) error {
+	return proto.UnmarshalMessageSetJSON(buf, m.XXX_extensions)
+}
 
 // ensure OldReply satisfies proto.Marshaler and proto.Unmarshaler
 var _ proto.Marshaler = (*OldReply)(nil)
diff --git a/protoc-gen-go/testdata/my_test/test.pb.go.golden b/protoc-gen-go/testdata/my_test/test.pb.go.golden
index 1165dea..361a813 100644
--- a/protoc-gen-go/testdata/my_test/test.pb.go.golden
+++ b/protoc-gen-go/testdata/my_test/test.pb.go.golden
@@ -413,6 +413,12 @@ func (m *OldReply) Marshal() ([]byte, error) {
 func (m *OldReply) Unmarshal(buf []byte) error {
 	return proto.UnmarshalMessageSet(buf, m.ExtensionMap())
 }
+func (m *OldReply) MarshalJSON() ([]byte, error) {
+	return proto.MarshalMessageSetJSON(m.XXX_extensions)
+}
+func (m *OldReply) UnmarshalJSON(buf []byte) error {
+	return proto.UnmarshalMessageSetJSON(buf, m.XXX_extensions)
+}
 
 // ensure OldReply satisfies proto.Marshaler and proto.Unmarshaler
 var _ proto.Marshaler = (*OldReply)(nil)