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>
diff --git a/encoding/protojson/encode.go b/encoding/protojson/encode.go
index fde1829..9b2592d 100644
--- a/encoding/protojson/encode.go
+++ b/encoding/protojson/encode.go
@@ -119,11 +119,21 @@
 
 	// 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 eef63d3..c7c3b8a 100644
--- a/encoding/prototext/encode.go
+++ b/encoding/prototext/encode.go
@@ -101,9 +101,15 @@
 	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
 		}