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>

diff --git a/src/reflect/all_test.go b/src/reflect/all_test.go
index df79f058075f4b58e19549ae24d9c039b39771f4..40ac6a95fa8db25b86c2e5e1a0736a892f8a300e 100644 (file)
--- 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 de01f13825fca1d5d9b624eef0c241f62bd065d2..a8274cc871f9ef3de9dbaacbccb58835cddb96bd 100644 (file)
--- 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