From fc5f8c340a1f4482e5f4ca4899ddd2cdbcdddeb7 Mon Sep 17 00:00:00 2001
From: Joe Tsai <joetsai@digital-static.net>
Date: Tue, 17 Sep 2019 22:32:53 -0700
Subject: [PATCH] encoding: optimize for oneofs

Suppose a oneof has N fields, the previous marshaling logic
would traverse every field checking for presence. This is O(N).
Using the protoreflect.Message.WhichOneof method, we can reduce
this to O(1). This optimization is exceptionally useful for oneofs
with a large number of fields.

Change-Id: I5f4aa8b1a899930f5c95e9cf1d68bac4b0b7884d
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/196121
Reviewed-by: Herbie Ong <herbie@google.com>
---
 encoding/protojson/encode.go | 14 ++++++++++++--
 encoding/prototext/encode.go | 10 ++++++++--
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/encoding/protojson/encode.go b/encoding/protojson/encode.go
index fde1829c..9b2592dd 100644
--- a/encoding/protojson/encode.go
+++ b/encoding/protojson/encode.go
@@ -119,11 +119,21 @@ func (o MarshalOptions) marshalFields(m pref.Message) error {
 
 	// Marshal out known fields.
 	fieldDescs := messageDesc.Fields()
-	for i := 0; i < fieldDescs.Len(); i++ {
+	for i := 0; i < fieldDescs.Len(); {
 		fd := fieldDescs.Get(i)
+		if od := fd.ContainingOneof(); od != nil {
+			fd = m.WhichOneof(od)
+			i += od.Fields().Len()
+			if fd == nil {
+				continue // unpopulated oneofs are not affected by EmitUnpopulated
+			}
+		} else {
+			i++
+		}
+
 		val := m.Get(fd)
 		if !m.Has(fd) {
-			if !o.EmitUnpopulated || fd.ContainingOneof() != nil {
+			if !o.EmitUnpopulated {
 				continue
 			}
 			isProto2Scalar := fd.Syntax() == pref.Proto2 && fd.Default().IsValid()
diff --git a/encoding/prototext/encode.go b/encoding/prototext/encode.go
index eef63d35..c7c3b8a4 100644
--- a/encoding/prototext/encode.go
+++ b/encoding/prototext/encode.go
@@ -101,9 +101,15 @@ func (o MarshalOptions) marshalMessage(m pref.Message) (text.Value, error) {
 	var msgFields [][2]text.Value
 	fieldDescs := messageDesc.Fields()
 	size := fieldDescs.Len()
-	for i := 0; i < size; i++ {
+	for i := 0; i < size; {
 		fd := fieldDescs.Get(i)
-		if !m.Has(fd) {
+		if od := fd.ContainingOneof(); od != nil {
+			fd = m.WhichOneof(od)
+			i += od.Fields().Len()
+		} else {
+			i++
+		}
+		if fd == nil || !m.Has(fd) {
 			continue
 		}