From 30c8da38eb656625d33d69d53d4b743a00d2b0d2 Mon Sep 17 00:00:00 2001
From: Josh Bleecher Snyder <josh@tailscale.com>
Date: Fri, 21 May 2021 09:43:58 -0700
Subject: [PATCH] reflect: add MapIter.Reset

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

Fixes #46293

Change-Id: I3aa6134dd00da35b508bd1e3b487332a871a3673
---
 src/reflect/all_test.go | 42 +++++++++++++++++++++++++++++++++++++++++
 src/reflect/value.go    | 12 ++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/src/reflect/all_test.go b/src/reflect/all_test.go
index 19b491c297d221..2602839d881f7d 100644
--- a/src/reflect/all_test.go
+++ b/src/reflect/all_test.go
@@ -7191,6 +7191,48 @@ 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.Fatal("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.Fatal("Next did not panic")
+	}()
+
+	// Reset to a diferent 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)
+	}
+
+	// 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 be159b8e7d9a05..f840bd64e90069 100644
--- a/src/reflect/value.go
+++ b/src/reflect/value.go
@@ -1357,6 +1357,18 @@ func (it *MapIter) Next() bool {
 	return mapiterkey(it.it) != nil
 }
 
+// Reset modifies it to iterate over v.
+// It panics if v's Kind is not Map and v is not the zero Value.
+func (it *MapIter) Reset(v Value) {
+	if v.IsValid() {
+		v.mustBe(Map)
+	}
+	it.m = v
+	it.it = nil
+	it.hiter = hiter{}
+	return
+}
+
 // MapRange returns a range iterator for a map.
 // It panics if v's Kind is not Map.
 //