From 34c3b587f666a6bb3dc1972532b2c6485abef4d7 Mon Sep 17 00:00:00 2001
From: Josh Bleecher Snyder <josh@tailscale.com>
Date: Thu, 29 Jul 2021 14:40:19 -0700
Subject: [PATCH] [tailscale1.17] reflect: allocate hiter as part of MapIter

This reduces the number of allocations per
reflect map iteration from two to one.

For #46293

(cherry picked from golang.org/cl/321889)

Change-Id: Ibcff5f42fc512e637b6e460bad4518e7ac83d4c3
---
 src/reflect/all_test.go |  7 +++----
 src/reflect/value.go    | 37 ++++++++++++++++++++++++++++++-------
 src/runtime/map.go      |  8 +++-----
 3 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/src/reflect/all_test.go b/src/reflect/all_test.go
index a3af10308e7f9e..231ba0691eda5a 100644
--- a/src/reflect/all_test.go
+++ b/src/reflect/all_test.go
@@ -369,10 +369,9 @@ func TestMapIterSet(t *testing.T) {
 			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)
+	// Making a *MapIter allocates. This should be the only allocation.
+	if got != 1 {
+		t.Errorf("wanted 1 alloc, got %d", got)
 	}
 }
 
diff --git a/src/reflect/value.go b/src/reflect/value.go
index ee98bdc3b0d0bf..0aeb9e8a20908f 100644
--- a/src/reflect/value.go
+++ b/src/reflect/value.go
@@ -1553,7 +1553,8 @@ func (v Value) MapKeys() []Value {
 	if m != nil {
 		mlen = maplen(m)
 	}
-	it := mapiterinit(v.typ, m)
+	it := unsafe.Pointer(new(hiter))
+	mapiterinit(v.typ, m, it)
 	a := make([]Value, mlen)
 	var i int
 	for i = 0; i < len(a); i++ {
@@ -1570,11 +1571,34 @@ func (v Value) MapKeys() []Value {
 	return a[:i]
 }
 
+// hiter's structure matches runtime.hiter's structure.
+// Having a clone here allows us to embed a map iterator
+// inside type MapIter so that MapIters can be re-used
+// without doing any allocations.
+type hiter struct {
+	key         unsafe.Pointer
+	elem        unsafe.Pointer
+	t           unsafe.Pointer
+	h           unsafe.Pointer
+	buckets     unsafe.Pointer
+	bptr        unsafe.Pointer
+	overflow    *[]unsafe.Pointer
+	oldoverflow *[]unsafe.Pointer
+	startBucket uintptr
+	offset      uint8
+	wrapped     bool
+	B           uint8
+	i           uint8
+	bucket      uintptr
+	checkBucket uintptr
+}
+
 // A MapIter is an iterator for ranging over a map.
 // See Value.MapRange.
 type MapIter struct {
-	m  Value
-	it unsafe.Pointer
+	m     Value
+	it    unsafe.Pointer // nil or points to hiter field
+	hiter hiter
 }
 
 // Key returns the key of the iterator's current map entry.
@@ -1660,7 +1684,8 @@ func (it *MapIter) SetValue(dst Value) {
 // calls to Key, Value, or Next will panic.
 func (it *MapIter) Next() bool {
 	if it.it == nil {
-		it.it = mapiterinit(it.m.typ, it.m.pointer())
+		it.it = unsafe.Pointer(&it.hiter)
+		mapiterinit(it.m.typ, it.m.pointer(), it.it)
 	} else {
 		if mapiterkey(it.it) == nil {
 			panic("MapIter.Next called on exhausted iterator")
@@ -3196,10 +3221,8 @@ func mapassign(t *rtype, m unsafe.Pointer, key, val unsafe.Pointer)
 //go:noescape
 func mapdelete(t *rtype, m unsafe.Pointer, key unsafe.Pointer)
 
-// m escapes into the return value, but the caller of mapiterinit
-// doesn't let the return value escape.
 //go:noescape
-func mapiterinit(t *rtype, m unsafe.Pointer) unsafe.Pointer
+func mapiterinit(t *rtype, m unsafe.Pointer, it unsafe.Pointer)
 
 //go:noescape
 func mapiterkey(it unsafe.Pointer) (key unsafe.Pointer)
diff --git a/src/runtime/map.go b/src/runtime/map.go
index 111db56b01a2fe..e92f6f13d6cb2e 100644
--- a/src/runtime/map.go
+++ b/src/runtime/map.go
@@ -159,8 +159,8 @@ type bmap struct {
 }
 
 // A hash iteration structure.
-// If you modify hiter, also change cmd/compile/internal/reflectdata/reflect.go to indicate
-// the layout of this structure.
+// If you modify hiter, also change cmd/compile/internal/reflectdata/reflect.go
+// and reflect/value.go to match the layout of this structure.
 type hiter struct {
 	key         unsafe.Pointer // Must be in first position.  Write nil to indicate iteration end (see cmd/compile/internal/walk/range.go).
 	elem        unsafe.Pointer // Must be in second position (see cmd/compile/internal/walk/range.go).
@@ -1335,10 +1335,8 @@ func reflect_mapdelete(t *maptype, h *hmap, key unsafe.Pointer) {
 }
 
 //go:linkname reflect_mapiterinit reflect.mapiterinit
-func reflect_mapiterinit(t *maptype, h *hmap) *hiter {
-	it := new(hiter)
+func reflect_mapiterinit(t *maptype, h *hmap, it *hiter) {
 	mapiterinit(t, h, it)
-	return it
 }
 
 //go:linkname reflect_mapiternext reflect.mapiternext