reflect: add MapIter.Reset

This allows callers to do (amortized) allocation-free iteration
over many maps.

Fixes #46293

Change-Id: I3aa6134dd00da35b508bd1e3b487332a871a3673
Reviewed-on: https://go-review.googlesource.com/c/go/+/321891
Trust: Josh Bleecher Snyder <josharian@gmail.com>
Run-TryBot: Josh Bleecher Snyder <josharian@gmail.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>
Reviewed-by: Keith Randall <khr@golang.org>

diff --git a/src/reflect/all_test.go b/src/reflect/all_test.go
index 6cb603cb162ba9db17a92c545574e0502e735607..01ce8b0c43cd6f2c7b13a6fdf285ccbae27a4488 100644 (file)
--- a/src/reflect/all_test.go
+++ b/src/reflect/all_test.go
@@ -7223,6 +7223,72 @@ func TestMapIterNilMap(t *testing.T) {
        }
 }
 
+func TestMapIterReset(t *testing.T) {
+       iter := new(MapIter)
+
+       // Use of zero iterator should panic.
+       func() {
+               defer func() { recover() }()
+               iter.Next()
+               t.Error("Next did not panic")
+       }()
+
+       // Reset to new Map should work.
+       m := map[string]int{"one": 1, "two": 2, "three": 3}
+       iter.Reset(ValueOf(m))
+       if got, want := iterateToString(iter), `[one: 1, three: 3, two: 2]`; got != want {
+               t.Errorf("iterator returned %s (after sorting), want %s", got, want)
+       }
+
+       // Reset to Zero value should work, but iterating over it should panic.
+       iter.Reset(Value{})
+       func() {
+               defer func() { recover() }()
+               iter.Next()
+               t.Error("Next did not panic")
+       }()
+
+       // Reset to a different Map with different types should work.
+       m2 := map[int]string{1: "one", 2: "two", 3: "three"}
+       iter.Reset(ValueOf(m2))
+       if got, want := iterateToString(iter), `[1: one, 2: two, 3: three]`; got != want {
+               t.Errorf("iterator returned %s (after sorting), want %s", got, want)
+       }
+
+       // Check that Reset, Next, and SetKey/SetValue play nicely together.
+       m3 := map[uint64]uint64{
+               1 << 0: 1 << 1,
+               1 << 1: 1 << 2,
+               1 << 2: 1 << 3,
+       }
+       kv := New(TypeOf(uint64(0))).Elem()
+       for i := 0; i < 5; i++ {
+               var seenk, seenv uint64
+               iter.Reset(ValueOf(m3))
+               for iter.Next() {
+                       iter.SetKey(kv)
+                       seenk ^= kv.Uint()
+                       iter.SetValue(kv)
+                       seenv ^= kv.Uint()
+               }
+               if seenk != 0b111 {
+                       t.Errorf("iteration yielded keys %b, want %b", seenk, 0b111)
+               }
+               if seenv != 0b1110 {
+                       t.Errorf("iteration yielded values %b, want %b", seenv, 0b1110)
+               }
+       }
+
+       // Reset should not allocate.
+       n := int(testing.AllocsPerRun(10, func() {
+               iter.Reset(ValueOf(m2))
+               iter.Reset(Value{})
+       }))
+       if n > 0 {
+               t.Errorf("MapIter.Reset allocated %d times", n)
+       }
+}
+
 func TestMapIterSafety(t *testing.T) {
        // Using a zero MapIter causes a panic, but not a crash.
        func() {

diff --git a/src/reflect/value.go b/src/reflect/value.go
index 90f31bae0a9a11d90825a1eb76bd5343a36eba4a..39c3c0f06e44359963ad69fc24b355c574f01591 100644 (file)
--- a/src/reflect/value.go
+++ b/src/reflect/value.go
@@ -1700,6 +1700,18 @@ func (it *MapIter) Next() bool {
        return mapiterkey(&it.hiter) != nil
 }
 
+// Reset modifies it to iterate over v.
+// It panics if v's Kind is not Map and v is not the zero Value.
+// Reset(Value{}) causes it to not to refer to any map,
+// which may allow the previously iterated-over map to be garbage collected.
+func (it *MapIter) Reset(v Value) {
+       if v.IsValid() {
+               v.mustBe(Map)
+       }
+       it.m = v
+       it.hiter = hiter{}
+}
+
 // MapRange returns a range iterator for a map.
 // It panics if v's Kind is not Map.
 //