[release-branch.go1.18] reflect: ensure map keys match key type in MapIndex and SetMapIndex

name                          old time/op    new time/op    delta
Map/StringKeys/MapIndex-8           2.36µs ± 5%    2.55µs ±11%  +7.98%  (p=0.006 n=10+9)
Map/StringKeys/SetMapIndex-8        4.86µs ± 7%    4.77µs ± 1%    ~     (p=0.211 n=10+9)
Map/StringKindKeys/MapIndex-8       2.29µs ± 3%    2.28µs ± 4%    ~     (p=0.631 n=10+10)
Map/StringKindKeys/SetMapIndex-8    4.44µs ± 3%    4.61µs ± 1%  +3.78%  (p=0.000 n=10+10)
Map/Uint64Keys/MapIndex-8           3.42µs ± 9%    3.11µs ± 2%  -9.20%  (p=0.000 n=10+9)
Map/Uint64Keys/SetMapIndex-8        5.17µs ± 3%    5.00µs ± 1%  -3.23%  (p=0.000 n=9+10)

For #52379
Fixes #52386

Change-Id: I545c71ea3145280828ca4186aad036a6c02016ed
Reviewed-on: https://go-review.googlesource.com/c/go/+/400635
Reviewed-by: Dmitri Shuralyov <dmitshur@google.com>
Reviewed-by: Joseph Tsai <joetsai@digital-static.net>
Reviewed-by: Ian Lance Taylor <iant@google.com>
Run-TryBot: Ian Lance Taylor <iant@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Auto-Submit: Ian Lance Taylor <iant@google.com>
(cherry picked from commit 11a650bb4aabfd7998b61df3ac33f61831d7abff)
Reviewed-on: https://go-review.googlesource.com/c/go/+/404000
Reviewed-by: David Chase <drchase@google.com>

2 files changed, 16 insertions, 2 deletions

diff --git a/src/reflect/all_test.go b/src/reflect/all_test.go
index 866f38e687..61663437d2 100644
--- a/src/reflect/all_test.go
+++ b/src/reflect/all_test.go
@@ -1496,6 +1496,10 @@ func TestMap(t *testing.T) {
 	if m != nil {
 		t.Errorf("mv.Set(nil) failed")
 	}
+
+	type S string
+	shouldPanic("not assignable", func() { mv.MapIndex(ValueOf(S("key"))) })
+	shouldPanic("not assignable", func() { mv.SetMapIndex(ValueOf(S("key")), ValueOf(0)) })
 }
 
 func TestNilMap(t *testing.T) {
@@ -7233,11 +7237,14 @@ func BenchmarkNew(b *testing.B) {
 
 func BenchmarkMap(b *testing.B) {
 	type V *int
+	type S string
 	value := ValueOf((V)(nil))
 	stringKeys := []string{}
 	mapOfStrings := map[string]V{}
 	uint64Keys := []uint64{}
 	mapOfUint64s := map[uint64]V{}
+	userStringKeys := []S{}
+	mapOfUserStrings := map[S]V{}
 	for i := 0; i < 100; i++ {
 		stringKey := fmt.Sprintf("key%d", i)
 		stringKeys = append(stringKeys, stringKey)
@@ -7246,6 +7253,10 @@ func BenchmarkMap(b *testing.B) {
 		uint64Key := uint64(i)
 		uint64Keys = append(uint64Keys, uint64Key)
 		mapOfUint64s[uint64Key] = nil
+
+		userStringKey := S(fmt.Sprintf("key%d", i))
+		userStringKeys = append(userStringKeys, userStringKey)
+		mapOfUserStrings[userStringKey] = nil
 	}
 
 	tests := []struct {
@@ -7254,6 +7265,7 @@ func BenchmarkMap(b *testing.B) {
 	}{
 		{"StringKeys", ValueOf(mapOfStrings), ValueOf(stringKeys), value},
 		{"Uint64Keys", ValueOf(mapOfUint64s), ValueOf(uint64Keys), value},
+		{"UserStringKeys", ValueOf(mapOfUserStrings), ValueOf(userStringKeys), value},
 	}
 
 	for _, tt := range tests {
diff --git a/src/reflect/value.go b/src/reflect/value.go
index dcc359dae4..147b402b2a 100644
--- a/src/reflect/value.go
+++ b/src/reflect/value.go
@@ -1583,6 +1583,8 @@ func (v Value) Len() int {
 	panic(&ValueError{"reflect.Value.Len", v.kind()})
 }
 
+var stringType = TypeOf("").(*rtype)
+
 // MapIndex returns the value associated with key in the map v.
 // It panics if v's Kind is not Map.
 // It returns the zero Value if key is not found in the map or if v represents a nil map.
@@ -1600,7 +1602,7 @@ func (v Value) MapIndex(key Value) Value {
 	// of unexported fields.
 
 	var e unsafe.Pointer
-	if key.kind() == String && tt.key.Kind() == String && tt.elem.size <= maxValSize {
+	if (tt.key == stringType || key.kind() == String) && tt.key == key.typ && tt.elem.size <= maxValSize {
 		k := *(*string)(key.ptr)
 		e = mapaccess_faststr(v.typ, v.pointer(), k)
 	} else {
@@ -2213,7 +2215,7 @@ func (v Value) SetMapIndex(key, elem Value) {
 	key.mustBeExported()
 	tt := (*mapType)(unsafe.Pointer(v.typ))
 
-	if key.kind() == String && tt.key.Kind() == String && tt.elem.size <= maxValSize {
+	if (tt.key == stringType || key.kind() == String) && tt.key == key.typ && tt.elem.size <= maxValSize {
 		k := *(*string)(key.ptr)
 		if elem.typ == nil {
 			mapdelete_faststr(v.typ, v.pointer(), k)