reflect: add MapIter.SetKey and MapIter.SetValue

These augment the existing MapIter.Key and MapIter.Value methods.
The existing methods return new Values.
Constructing these new Values often requires allocating.
These methods allow the caller to bring their own storage.

The naming is somewhat unfortunate, in that the spec
uses the word "element" instead of "value",
as do the reflect.Type methods.
In a vacuum, MapIter.SetElem would be preferable.
However, matching the existing methods is more important.

Fixes #32424
Fixes #46131

Change-Id: I19c4d95c432f63dfe52cde96d2125abd021f24fa
Reviewed-on: https://go-review.googlesource.com/c/go/+/320929
Trust: Josh Bleecher Snyder <josharian@gmail.com>
Run-TryBot: Josh Bleecher Snyder <josharian@gmail.com>
Reviewed-by: Keith Randall <khr@golang.org>

2 files changed, 99 insertions, 4 deletions

diff --git a/src/reflect/all_test.go b/src/reflect/all_test.go
index df79f05807..40ac6a95fa 100644
--- a/src/reflect/all_test.go
+++ b/src/reflect/all_test.go
@@ -336,6 +336,47 @@ func TestSetValue(t *testing.T) {
 	}
 }
 
+func TestMapIterSet(t *testing.T) {
+	m := make(map[string]interface{}, len(valueTests))
+	for _, tt := range valueTests {
+		m[tt.s] = tt.i
+	}
+	v := ValueOf(m)
+
+	k := New(v.Type().Key()).Elem()
+	e := New(v.Type().Elem()).Elem()
+
+	iter := v.MapRange()
+	for iter.Next() {
+		iter.SetKey(k)
+		iter.SetValue(e)
+		want := m[k.String()]
+		got := e.Interface()
+		if got != want {
+			t.Errorf("%q: want (%T) %v, got (%T) %v", k.String(), want, want, got, got)
+		}
+		if setkey, key := valueToString(k), valueToString(iter.Key()); setkey != key {
+			t.Errorf("MapIter.Key() = %q, MapIter.SetKey() = %q", key, setkey)
+		}
+		if setval, val := valueToString(e), valueToString(iter.Value()); setval != val {
+			t.Errorf("MapIter.Value() = %q, MapIter.SetValue() = %q", val, setval)
+		}
+	}
+
+	got := int(testing.AllocsPerRun(10, func() {
+		iter := v.MapRange()
+		for iter.Next() {
+			iter.SetKey(k)
+			iter.SetValue(e)
+		}
+	}))
+	// Making a *MapIter and making an hiter both allocate.
+	// Those should be the only two allocations.
+	if got != 2 {
+		t.Errorf("wanted 2 allocs, got %d", got)
+	}
+}
+
 func TestCanSetField(t *testing.T) {
 	type embed struct{ x, X int }
 	type Embed struct{ x, X int }
diff --git a/src/reflect/value.go b/src/reflect/value.go
index de01f13825..a8274cc871 100644
--- a/src/reflect/value.go
+++ b/src/reflect/value.go
@@ -1578,13 +1578,40 @@ func (it *MapIter) Key() Value {
 	if it.it == nil {
 		panic("MapIter.Key called before Next")
 	}
-	if mapiterkey(it.it) == nil {
+	iterkey := mapiterkey(it.it)
+	if iterkey == nil {
 		panic("MapIter.Key called on exhausted iterator")
 	}
 
 	t := (*mapType)(unsafe.Pointer(it.m.typ))
 	ktype := t.key
-	return copyVal(ktype, it.m.flag.ro()|flag(ktype.Kind()), mapiterkey(it.it))
+	return copyVal(ktype, it.m.flag.ro()|flag(ktype.Kind()), iterkey)
+}
+
+// SetKey assigns dst to the key of the iterator's current map entry.
+// It is equivalent to dst.Set(it.Key()), but it avoids allocating a new Value.
+// As in Go, the key must be assignable to dst's type.
+func (it *MapIter) SetKey(dst Value) {
+	if it.it == nil {
+		panic("MapIter.SetKey called before Next")
+	}
+	iterkey := mapiterkey(it.it)
+	if iterkey == nil {
+		panic("MapIter.SetKey called on exhausted iterator")
+	}
+
+	dst.mustBeAssignable()
+	var target unsafe.Pointer
+	if dst.kind() == Interface {
+		target = dst.ptr
+	}
+
+	t := (*mapType)(unsafe.Pointer(it.m.typ))
+	ktype := t.key
+
+	key := Value{ktype, iterkey, it.m.flag | flag(ktype.Kind())}
+	key = key.assignTo("reflect.MapIter.SetKey", dst.typ, target)
+	typedmemmove(dst.typ, dst.ptr, key.ptr)
 }
 
 // Value returns the value of the iterator's current map entry.
@@ -1592,13 +1619,40 @@ func (it *MapIter) Value() Value {
 	if it.it == nil {
 		panic("MapIter.Value called before Next")
 	}
-	if mapiterkey(it.it) == nil {
+	iterelem := mapiterelem(it.it)
+	if iterelem == nil {
 		panic("MapIter.Value called on exhausted iterator")
 	}
 
 	t := (*mapType)(unsafe.Pointer(it.m.typ))
 	vtype := t.elem
-	return copyVal(vtype, it.m.flag.ro()|flag(vtype.Kind()), mapiterelem(it.it))
+	return copyVal(vtype, it.m.flag.ro()|flag(vtype.Kind()), iterelem)
+}
+
+// SetValue assigns dst to the value of the iterator's current map entry.
+// It is equivalent to dst.Set(it.Value()), but it avoids allocating a new Value.
+// As in Go, the value must be assignable to dst's type.
+func (it *MapIter) SetValue(dst Value) {
+	if it.it == nil {
+		panic("MapIter.SetValue called before Next")
+	}
+	iterelem := mapiterelem(it.it)
+	if iterelem == nil {
+		panic("MapIter.SetValue called on exhausted iterator")
+	}
+
+	dst.mustBeAssignable()
+	var target unsafe.Pointer
+	if dst.kind() == Interface {
+		target = dst.ptr
+	}
+
+	t := (*mapType)(unsafe.Pointer(it.m.typ))
+	vtype := t.elem
+
+	elem := Value{vtype, iterelem, it.m.flag | flag(vtype.Kind())}
+	elem = elem.assignTo("reflect.MapIter.SetValue", dst.typ, target)
+	typedmemmove(dst.typ, dst.ptr, elem.ptr)
 }
 
 // Next advances the map iterator and reports whether there is another