reflect: allocate hiter as part of MapIter

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

For #46293

Change-Id: Ibcff5f42fc512e637b6e460bad4518e7ac83d4c3
Reviewed-on: https://go-review.googlesource.com/c/go/+/321889
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>

src/reflect/all_test.go

@@ -370,10 +370,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)
 	}
 }
 

src/reflect/value.go

@@ -1549,11 +1549,12 @@ func (v Value) MapKeys() []Value {
 	if m != nil {
 		mlen = maplen(m)
 	}
-	it := mapiterinit(v.typ, m)
+	var it hiter
+	mapiterinit(v.typ, m, &it)
 	a := make([]Value, mlen)
 	var i int
 	for i = 0; i < len(a); i++ {
-		key := mapiterkey(it)
+		key := mapiterkey(&it)
 		if key == nil {
 			// Someone deleted an entry from the map since we
 			// called maplen above. It's a data race, but nothing
@@ -1561,24 +1562,50 @@ func (v Value) MapKeys() []Value {
 			break
 		}
 		a[i] = copyVal(keyType, fl, key)
-		mapiternext(it)
+		mapiternext(&it)
 	}
 	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
+}
+
+func (h hiter) initialized() bool {
+	return h.t != nil
+}
+
 // A MapIter is an iterator for ranging over a map.
 // See Value.MapRange.
 type MapIter struct {
-	m  Value
-	it unsafe.Pointer
+	m     Value
+	hiter hiter
 }
 
 // Key returns the key of the iterator's current map entry.
 func (it *MapIter) Key() Value {
-	if it.it == nil {
+	if !it.hiter.initialized() {
 		panic("MapIter.Key called before Next")
 	}
-	iterkey := mapiterkey(it.it)
+	iterkey := mapiterkey(&it.hiter)
 	if iterkey == nil {
 		panic("MapIter.Key called on exhausted iterator")
 	}
@@ -1592,10 +1619,10 @@ func (it *MapIter) Key() Value {
 // It is equivalent to dst.Set(it.Key()), but it avoids allocating a new Value.
 // As in Go, the key must be assignable to dst's type.
 func (it *MapIter) SetKey(dst Value) {
-	if it.it == nil {
+	if !it.hiter.initialized() {
 		panic("MapIter.SetKey called before Next")
 	}
-	iterkey := mapiterkey(it.it)
+	iterkey := mapiterkey(&it.hiter)
 	if iterkey == nil {
 		panic("MapIter.SetKey called on exhausted iterator")
 	}
@@ -1616,10 +1643,10 @@ func (it *MapIter) SetKey(dst Value) {
 
 // Value returns the value of the iterator's current map entry.
 func (it *MapIter) Value() Value {
-	if it.it == nil {
+	if !it.hiter.initialized() {
 		panic("MapIter.Value called before Next")
 	}
-	iterelem := mapiterelem(it.it)
+	iterelem := mapiterelem(&it.hiter)
 	if iterelem == nil {
 		panic("MapIter.Value called on exhausted iterator")
 	}
@@ -1633,10 +1660,10 @@ func (it *MapIter) Value() Value {
 // It is equivalent to dst.Set(it.Value()), but it avoids allocating a new Value.
 // As in Go, the value must be assignable to dst's type.
 func (it *MapIter) SetValue(dst Value) {
-	if it.it == nil {
+	if !it.hiter.initialized() {
 		panic("MapIter.SetValue called before Next")
 	}
-	iterelem := mapiterelem(it.it)
+	iterelem := mapiterelem(&it.hiter)
 	if iterelem == nil {
 		panic("MapIter.SetValue called on exhausted iterator")
 	}
@@ -1659,15 +1686,15 @@ func (it *MapIter) SetValue(dst Value) {
 // entry. It returns false when the iterator is exhausted; subsequent
 // 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())
+	if !it.hiter.initialized() {
+		mapiterinit(it.m.typ, it.m.pointer(), &it.hiter)
 	} else {
-		if mapiterkey(it.it) == nil {
+		if mapiterkey(&it.hiter) == nil {
 			panic("MapIter.Next called on exhausted iterator")
 		}
-		mapiternext(it.it)
+		mapiternext(&it.hiter)
 	}
-	return mapiterkey(it.it) != nil
+	return mapiterkey(&it.hiter) != nil
 }
 
 // MapRange returns a range iterator for a map.
@@ -3216,19 +3243,17 @@ 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 *hiter)
 
 //go:noescape
-func mapiterkey(it unsafe.Pointer) (key unsafe.Pointer)
+func mapiterkey(it *hiter) (key unsafe.Pointer)
 
 //go:noescape
-func mapiterelem(it unsafe.Pointer) (elem unsafe.Pointer)
+func mapiterelem(it *hiter) (elem unsafe.Pointer)
 
 //go:noescape
-func mapiternext(it unsafe.Pointer)
+func mapiternext(it *hiter)
 
 //go:noescape
 func maplen(m unsafe.Pointer) int

src/runtime/map.go

@@ -160,8 +160,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).
@@ -806,6 +806,7 @@ func mapiterinit(t *maptype, h *hmap, it *hiter) {
 		racereadpc(unsafe.Pointer(h), callerpc, abi.FuncPCABIInternal(mapiterinit))
 	}
 
+	it.t = t
 	if h == nil || h.count == 0 {
 		return
 	}
@@ -813,7 +814,6 @@ func mapiterinit(t *maptype, h *hmap, it *hiter) {
 	if unsafe.Sizeof(hiter{})/goarch.PtrSize != 12 {
 		throw("hash_iter size incorrect") // see cmd/compile/internal/reflectdata/reflect.go
 	}
-	it.t = t
 	it.h = h
 
 	// grab snapshot of bucket state
@@ -1336,10 +1336,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