diff --git a/src/reflect/all_test.go b/src/reflect/all_test.go index 6cb603cb162..01ce8b0c43c 100644 --- 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 90f31bae0a9..39c3c0f06e4 100644 --- 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. //