From ca0b25e48ff647e285827e8af0be93f82fa144a6 Mon Sep 17 00:00:00 2001
From: Joe Tsai <joetsai@digital-static.net>
Date: Thu, 27 Feb 2020 14:27:08 -0800
Subject: [PATCH] testing/protocmp: automatically promote message values

For user convenience, automatically transform message values by
shallow copying them if necessary. Storing messages as values is
frowned upon, but is sometimes done by APIs that a user does not own.

Change-Id: I7e927d1a1e050bf4cea1aa889f56d23e99355f26
Reviewed-on: https://go-review.googlesource.com/c/protobuf/+/221423
Reviewed-by: Damien Neil <dneil@google.com>
---
 testing/protocmp/util_test.go | 59 +++++++++++++++++++++++++++++++++++
 testing/protocmp/xform.go     | 20 ++++++++++--
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/testing/protocmp/util_test.go b/testing/protocmp/util_test.go
index f9da7e52..6f41ba4f 100644
--- a/testing/protocmp/util_test.go
+++ b/testing/protocmp/util_test.go
@@ -145,6 +145,65 @@ func TestEqual(t *testing.T) {
 		want: true,
 	}}...)
 
+	// Test message values.
+	tests = append(tests, []test{{
+		x:    testpb.TestAllTypes{OptionalSint64: proto.Int64(1)},
+		y:    testpb.TestAllTypes{OptionalSint64: proto.Int64(1)},
+		opts: cmp.Options{Transform()},
+		want: true,
+	}, {
+		x:    testpb.TestAllTypes{OptionalSint64: proto.Int64(1)},
+		y:    testpb.TestAllTypes{OptionalSint64: proto.Int64(2)},
+		opts: cmp.Options{Transform()},
+		want: false,
+	}, {
+		x:    struct{ M testpb.TestAllTypes }{M: testpb.TestAllTypes{OptionalSint64: proto.Int64(1)}},
+		y:    struct{ M testpb.TestAllTypes }{M: testpb.TestAllTypes{OptionalSint64: proto.Int64(1)}},
+		opts: cmp.Options{Transform()},
+		want: true,
+	}, {
+		x:    struct{ M testpb.TestAllTypes }{M: testpb.TestAllTypes{OptionalSint64: proto.Int64(1)}},
+		y:    struct{ M testpb.TestAllTypes }{M: testpb.TestAllTypes{OptionalSint64: proto.Int64(2)}},
+		opts: cmp.Options{Transform()},
+		want: false,
+	}, {
+		x:    struct{ M []testpb.TestAllTypes }{M: []testpb.TestAllTypes{{OptionalSint64: proto.Int64(1)}}},
+		y:    struct{ M []testpb.TestAllTypes }{M: []testpb.TestAllTypes{{OptionalSint64: proto.Int64(1)}}},
+		opts: cmp.Options{Transform()},
+		want: true,
+	}, {
+		x:    struct{ M []testpb.TestAllTypes }{M: []testpb.TestAllTypes{{OptionalSint64: proto.Int64(1)}}},
+		y:    struct{ M []testpb.TestAllTypes }{M: []testpb.TestAllTypes{{OptionalSint64: proto.Int64(2)}}},
+		opts: cmp.Options{Transform()},
+		want: false,
+	}, {
+		x: struct {
+			M map[string]testpb.TestAllTypes
+		}{
+			M: map[string]testpb.TestAllTypes{"k": {OptionalSint64: proto.Int64(1)}},
+		},
+		y: struct {
+			M map[string]testpb.TestAllTypes
+		}{
+			M: map[string]testpb.TestAllTypes{"k": {OptionalSint64: proto.Int64(1)}},
+		},
+		opts: cmp.Options{Transform()},
+		want: true,
+	}, {
+		x: struct {
+			M map[string]testpb.TestAllTypes
+		}{
+			M: map[string]testpb.TestAllTypes{"k": {OptionalSint64: proto.Int64(1)}},
+		},
+		y: struct {
+			M map[string]testpb.TestAllTypes
+		}{
+			M: map[string]testpb.TestAllTypes{"k": {OptionalSint64: proto.Int64(2)}},
+		},
+		opts: cmp.Options{Transform()},
+		want: false,
+	}}...)
+
 	// Test IgnoreUnknown.
 	raw := pack.Message{
 		pack.Tag{1, pack.BytesType}, pack.String("Hello, goodbye!"),
diff --git a/testing/protocmp/xform.go b/testing/protocmp/xform.go
index e27c8908..a42bf9bb 100644
--- a/testing/protocmp/xform.go
+++ b/testing/protocmp/xform.go
@@ -161,10 +161,18 @@ func Transform(...option) cmp.Option {
 	// NOTE: There are currently no custom options for Transform,
 	// but the use of an unexported type keeps the future open.
 
+	// addrType returns a pointer to t if t isn't a pointer or interface.
+	addrType := func(t reflect.Type) reflect.Type {
+		if k := t.Kind(); k == reflect.Interface || k == reflect.Ptr {
+			return t
+		}
+		return reflect.PtrTo(t)
+	}
+
 	// TODO: Should this transform protoreflect.Enum types to Enum as well?
 	return cmp.FilterPath(func(p cmp.Path) bool {
 		ps := p.Last()
-		if isMessageType(ps.Type()) {
+		if isMessageType(addrType(ps.Type())) {
 			return true
 		}
 
@@ -175,11 +183,19 @@ func Transform(...option) cmp.Option {
 			if !vx.IsValid() || vx.IsNil() || !vy.IsValid() || vy.IsNil() {
 				return false
 			}
-			return isMessageType(vx.Elem().Type()) && isMessageType(vy.Elem().Type())
+			return isMessageType(addrType(vx.Elem().Type())) && isMessageType(addrType(vy.Elem().Type()))
 		}
 
 		return false
 	}, cmp.Transformer("protocmp.Transform", func(v interface{}) Message {
+		// For user convenience, shallow copy the message value if necessary
+		// in order for it to implement the message interface.
+		if rv := reflect.ValueOf(v); rv.IsValid() && rv.Kind() != reflect.Ptr && !isMessageType(rv.Type()) {
+			pv := reflect.New(rv.Type())
+			pv.Elem().Set(rv)
+			v = pv.Interface()
+		}
+
 		m := protoimpl.X.MessageOf(v)
 		switch {
 		case m == nil: