internal/runtime/maps: store group across Iter.Next calls

A previous CL kept it across loop iterations, but those are more rare
than call iterations.

For #54766.

Cq-Include-Trybots: luci.golang.try:gotip-linux-amd64-longtest-swissmap
Change-Id: Ieea0f1677e357f5e451650b1c697da7f63f3bca1
Reviewed-on: https://go-review.googlesource.com/c/go/+/621116
Reviewed-by: Keith Randall <khr@google.com>
Auto-Submit: Michael Pratt <mpratt@google.com>
Reviewed-by: Keith Randall <khr@golang.org>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

1 files changed, 16 insertions, 16 deletions

diff --git a/src/internal/runtime/maps/table.go b/src/internal/runtime/maps/table.go
index bda74ea41b..d5ec24e7a6 100644
--- a/src/internal/runtime/maps/table.go
+++ b/src/internal/runtime/maps/table.go
@@ -540,8 +540,10 @@ type Iter struct {
 	dirIdx int
 
 	// tab is the table at dirIdx during the previous call to Next.
-	tab        *table
-	groupSmall groupReference // only if small map at init
+	tab *table
+
+	// group is the group at entryIdx during the previous call to Next.
+	group groupReference
 
 	// entryIdx is the current entry index, prior to adjustment by entryOffset.
 	// The lower 3 bits of the index are the slot index, and the upper bits
@@ -570,7 +572,7 @@ func (it *Iter) Init(typ *abi.SwissMapType, m *Map) {
 	it.dirOffset = rand()
 	it.globalDepth = m.globalDepth
 	it.dirIdx = dirIdx
-	it.groupSmall = groupSmall
+	it.group = groupSmall
 	it.clearSeq = m.clearSeq
 }
 
@@ -619,16 +621,15 @@ func (it *Iter) Next() {
 
 	if it.dirIdx < 0 {
 		// Map was small at Init.
-		g := it.groupSmall
 		for ; it.entryIdx < abi.SwissMapGroupSlots; it.entryIdx++ {
 			k := uintptr(it.entryIdx+it.entryOffset) % abi.SwissMapGroupSlots
 
-			if (g.ctrls().get(k) & ctrlEmpty) == ctrlEmpty {
+			if (it.group.ctrls().get(k) & ctrlEmpty) == ctrlEmpty {
 				// Empty or deleted.
 				continue
 			}
 
-			key := g.key(it.typ, k)
+			key := it.group.key(it.typ, k)
 			if it.typ.IndirectKey() {
 				key = *((*unsafe.Pointer)(key))
 			}
@@ -645,7 +646,7 @@ func (it *Iter) Next() {
 				if !ok {
 					// See comment below.
 					if it.clearSeq == it.m.clearSeq && !it.typ.Key.Equal(key, key) {
-						elem = g.elem(it.typ, k)
+						elem = it.group.elem(it.typ, k)
 						if it.typ.IndirectElem() {
 							elem = *((*unsafe.Pointer)(elem))
 						}
@@ -657,7 +658,7 @@ func (it *Iter) Next() {
 					elem = newElem
 				}
 			} else {
-				elem = g.elem(it.typ, k)
+				elem = it.group.elem(it.typ, k)
 				if it.typ.IndirectElem() {
 					elem = *((*unsafe.Pointer)(elem))
 				}
@@ -736,8 +737,6 @@ func (it *Iter) Next() {
 			it.tab = newTab
 		}
 
-		var g groupReference
-
 		// N.B. Use it.tab, not newTab. It is important to use the old
 		// table for key selection if the table has grown. See comment
 		// on grown below.
@@ -745,25 +744,25 @@ func (it *Iter) Next() {
 			entryIdx := (it.entryIdx + it.entryOffset) & it.tab.groups.entryMask
 			slotIdx := uintptr(entryIdx & (abi.SwissMapGroupSlots - 1))
 
-			if slotIdx == 0 || g.data == nil {
+			if slotIdx == 0 || it.group.data == nil {
 				// Only compute the group (a) when we switch
 				// groups (slotIdx rolls over) and (b) on the
 				// first iteration in this table (slotIdx may
 				// not be zero due to entryOffset).
 				groupIdx := entryIdx >> abi.SwissMapGroupSlotsBits
-				g = it.tab.groups.group(it.typ, groupIdx)
+				it.group = it.tab.groups.group(it.typ, groupIdx)
 			}
 
 			// TODO(prattmic): Skip over groups that are composed of only empty
 			// or deleted slots using matchEmptyOrDeleted() and counting the
 			// number of bits set.
 
-			if (g.ctrls().get(slotIdx) & ctrlEmpty) == ctrlEmpty {
+			if (it.group.ctrls().get(slotIdx) & ctrlEmpty) == ctrlEmpty {
 				// Empty or deleted.
 				continue
 			}
 
-			key := g.key(it.typ, slotIdx)
+			key := it.group.key(it.typ, slotIdx)
 			if it.typ.IndirectKey() {
 				key = *((*unsafe.Pointer)(key))
 			}
@@ -809,7 +808,7 @@ func (it *Iter) Next() {
 					// need to return anything added after
 					// clear.
 					if it.clearSeq == it.m.clearSeq && !it.typ.Key.Equal(key, key) {
-						elem = g.elem(it.typ, slotIdx)
+						elem = it.group.elem(it.typ, slotIdx)
 						if it.typ.IndirectElem() {
 							elem = *((*unsafe.Pointer)(elem))
 						}
@@ -821,7 +820,7 @@ func (it *Iter) Next() {
 					elem = newElem
 				}
 			} else {
-				elem = g.elem(it.typ, slotIdx)
+				elem = it.group.elem(it.typ, slotIdx)
 				if it.typ.IndirectElem() {
 					elem = *((*unsafe.Pointer)(elem))
 				}
@@ -862,6 +861,7 @@ func (it *Iter) Next() {
 		entries := 1 << (it.m.globalDepth - it.tab.localDepth)
 		it.dirIdx += entries
 		it.tab = nil
+		it.group = groupReference{}
 		it.entryIdx = 0
 	}