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 <[email protected]>
Run-TryBot: Josh Bleecher Snyder <[email protected]>
TryBot-Result: Go Bot <[email protected]>
Reviewed-by: Ian Lance Taylor <[email protected]>
Reviewed-by: Keith Randall <[email protected]>

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() {

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.
 //